tractor.dev/toolkit-go@v0.0.0-20241010005851-214d91207d07/duplex/mux/listen_ws.go (about)

     1  //go:build !tinygo
     2  
     3  package mux
     4  
     5  import (
     6  	"io"
     7  	"net"
     8  	"net/http"
     9  
    10  	"golang.org/x/net/websocket"
    11  )
    12  
    13  // wsListener wraps a net.Listener and WebSocket server to return connected mux sessions.
    14  type wsListener struct {
    15  	net.Listener
    16  	accepted chan Session
    17  }
    18  
    19  // Accept waits for and returns the next connected session to the listener.
    20  func (l *wsListener) Accept() (Session, error) {
    21  	sess, ok := <-l.accepted
    22  	if !ok {
    23  		return nil, io.EOF
    24  	}
    25  	return sess, nil
    26  }
    27  
    28  // Close closes the listener.
    29  // Any blocked Accept operations will be unblocked and return errors.
    30  func (l *wsListener) Close() error {
    31  	close(l.accepted)
    32  	return l.Listener.Close()
    33  }
    34  
    35  func (l *wsListener) Addr() net.Addr {
    36  	return l.Listener.Addr()
    37  }
    38  
    39  // ListenWS takes a TCP address and returns a Listener for a HTTP+WebSocket server listening on the given address.
    40  func ListenWS(addr string) (Listener, error) {
    41  	l, err := net.Listen("tcp", addr)
    42  	if err != nil {
    43  		return nil, err
    44  	}
    45  	wsl := &wsListener{
    46  		Listener: l,
    47  		accepted: make(chan Session),
    48  	}
    49  	srv := &http.Server{
    50  		Addr: addr,
    51  		Handler: websocket.Handler(func(ws *websocket.Conn) {
    52  			ws.PayloadType = websocket.BinaryFrame
    53  			sess := New(ws)
    54  			defer sess.Close()
    55  			wsl.accepted <- sess
    56  			sess.Wait()
    57  		}),
    58  	}
    59  	go srv.Serve(l)
    60  	return wsl, nil
    61  }