github.com/golazy/golazy@v0.0.7-0.20221012133820-968fe65a0b65/lazydev/protocolmux/protocolmux.go (about)

     1  package protocolmux
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"io"
     7  	"log"
     8  	"net"
     9  )
    10  
    11  type handler struct {
    12  	prefixes [][]byte
    13  	conns    chan (*conn)
    14  	mux      *Mux
    15  }
    16  
    17  func (h *handler) Accept() (net.Conn, error) {
    18  	c, ok := <-h.conns
    19  	if !ok {
    20  		return nil, io.EOF
    21  	}
    22  	return c, nil
    23  }
    24  
    25  func (h *handler) Addr() net.Addr {
    26  	return h.mux.L.Addr()
    27  }
    28  
    29  func (h *handler) Close() error {
    30  	for i, handler := range h.mux.handlers {
    31  		if handler == h {
    32  			h.mux.handlers = append(h.mux.handlers[:i], h.mux.handlers[i+1:]...)
    33  		}
    34  	}
    35  
    36  	close(h.conns)
    37  
    38  	return nil
    39  }
    40  
    41  type Mux struct {
    42  	L        net.Listener
    43  	handlers []*handler
    44  }
    45  
    46  type conn struct {
    47  	net.Conn
    48  	prefix []byte
    49  }
    50  
    51  func (m *Mux) ListenTo(prefixes [][]byte) net.Listener {
    52  	h := &handler{
    53  		prefixes: prefixes,
    54  		conns:    make(chan (*conn)),
    55  		mux:      m,
    56  	}
    57  	m.handlers = append(m.handlers, h)
    58  
    59  	return h
    60  }
    61  
    62  func (m *Mux) l(v ...interface{}) {
    63  	fmt.Println(append([]interface{}{"MUX:"}, v...)...)
    64  }
    65  
    66  func (m *Mux) Listen() error {
    67  	m.l("starting listener")
    68  	if m.L == nil {
    69  		return fmt.Errorf("listener not defined")
    70  	}
    71  
    72  	for {
    73  		m.l("Waiting for connection")
    74  		conn, err := m.L.Accept()
    75  		m.l("Got new Connection", conn.RemoteAddr())
    76  		fmt.Println()
    77  		if err != nil {
    78  			return err
    79  		}
    80  		go m.handleConn(conn)
    81  	}
    82  }
    83  
    84  func (c *conn) Read(b []byte) (int, error) {
    85  	if len(c.prefix) == 0 {
    86  		return c.Conn.Read(b)
    87  	}
    88  	n := copy(b, c.prefix)
    89  	c.prefix = c.prefix[n:] // BUG?
    90  	return n, nil
    91  }
    92  
    93  func (m *Mux) handleConn(c net.Conn) {
    94  	buf := make([]byte, 1024)
    95  	n, err := c.Read(buf)
    96  	if err != nil {
    97  		log.Println("erro while handling a connection:", err)
    98  	}
    99  
   100  	for _, handler := range m.handlers {
   101  		for _, prefix := range handler.prefixes {
   102  			if bytes.HasPrefix(buf[:n], prefix) {
   103  				c := &conn{Conn: c, prefix: buf[:n]}
   104  				handler.conns <- c
   105  				return
   106  			}
   107  		}
   108  	}
   109  	log.Println("no protocol for connection")
   110  	c.Close()
   111  }