github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/network/web_listener.go (about)

     1  package network
     2  
     3  import (
     4  	"github.com/15mga/kiwi"
     5  	"github.com/15mga/kiwi/util"
     6  	"github.com/fasthttp/websocket"
     7  	"net/http"
     8  )
     9  
    10  type webOption struct {
    11  	addr      string
    12  	upgrader  *websocket.Upgrader
    13  	resHeader http.Header
    14  }
    15  
    16  type WebOption func(option *webOption)
    17  
    18  func WebAddr(addr string) WebOption {
    19  	return func(option *webOption) {
    20  		option.addr = addr
    21  	}
    22  }
    23  
    24  func WebUpgrader(fn func(*websocket.Upgrader)) WebOption {
    25  	return func(option *webOption) {
    26  		fn(option.upgrader)
    27  	}
    28  }
    29  
    30  func WebResHeader(header http.Header) WebOption {
    31  	return func(option *webOption) {
    32  		option.resHeader = header
    33  	}
    34  }
    35  
    36  func NewWebListener(onConn func(conn *websocket.Conn), opts ...WebOption) kiwi.IListener {
    37  	o := &webOption{
    38  		addr:     ":7737",
    39  		upgrader: &websocket.Upgrader{},
    40  	}
    41  	for _, opt := range opts {
    42  		opt(o)
    43  	}
    44  	return &webListener{
    45  		option: o,
    46  		onConn: onConn,
    47  	}
    48  }
    49  
    50  type webListener struct {
    51  	server *http.Server
    52  	option *webOption
    53  	onConn func(conn *websocket.Conn)
    54  }
    55  
    56  func (l *webListener) Addr() string {
    57  	return l.option.addr
    58  }
    59  
    60  func (l *webListener) Port() int {
    61  	port, _ := util.ParseAddrPort(l.server.Addr)
    62  	return port
    63  }
    64  
    65  func (l *webListener) handler(writer http.ResponseWriter, request *http.Request) {
    66  	conn, e := l.option.upgrader.Upgrade(writer, request, l.option.resHeader)
    67  	if e != nil {
    68  		kiwi.Error3(util.EcServiceErr, e)
    69  		return
    70  	}
    71  	l.onConn(conn)
    72  }
    73  
    74  func (l *webListener) Start() *util.Err {
    75  	kiwi.Info("start websocket listener", util.M{
    76  		"addr": l.option.addr,
    77  	})
    78  	http.HandleFunc("/", l.handler)
    79  
    80  	go func() {
    81  		e := http.ListenAndServe(l.option.addr, nil)
    82  		if e != nil {
    83  			panic(e)
    84  		}
    85  	}()
    86  	return nil
    87  }
    88  
    89  func (l *webListener) Close() {
    90  	l.server.Close()
    91  }