add UpgradeV2 (#29)

antlabs · Jun 22, 2024 · fc60bda · fc60bda
1 parent 6eaa970
commit fc60bda
Show file tree

Hide file tree

Showing 8 changed files with 63 additions and 21 deletions.
diff --git a/autobahn/config/fuzzingclient.json b/autobahn/config/fuzzingclient.json
@@ -1,6 +1,13 @@
 {
  "outdir": "./report/",
  "servers": [
+ {
+ "agent": "global",
+ "url": "ws://localhost:9001/global",
+ "options": {
+ "version": 18
+ }
+ },
  {
  "agent": "no-context-takeover-decompression-and-compression-no-tls",
  "url": "ws://localhost:9001/no-context-takeover-decompression-and-compression",
@@ -37,4 +44,4 @@
  ""
  ],
  "exclude-agent-cases": {}
-}
+}
diff --git a/autobahn/server/autobahn-server.go b/autobahn/server/autobahn-server.go
@@ -136,6 +136,24 @@ func echoReadTime(w http.ResponseWriter, r *http.Request) {
  _ = c.ReadLoop()
 }
 
+var upgrade = quickws.NewUpgrade(
+ quickws.WithServerReplyPing(),
+ quickws.WithServerDecompression(),
+ quickws.WithServerIgnorePong(),
+ quickws.WithServerEnableUTF8Check(),
+ quickws.WithServerReadTimeout(5*time.Second),
+)
+
+func global(w http.ResponseWriter, r *http.Request) {
+ c, err := upgrade.UpgradeV2(w, r, &echoHandler{openWriteTimeout: true})
+ if err != nil {
+ fmt.Println("Upgrade fail:", err)
+ return
+ }
+
+ _ = c.ReadLoop()
+}
+
 func startTLSServer(mux *http.ServeMux) {
 
  cert, err := tls.X509KeyPair(certPEMBlock, keyPEMBlock)
@@ -167,6 +185,7 @@ func startServer(mux *http.ServeMux) {
 func main() {
  mux := &http.ServeMux{}
  mux.HandleFunc("/timeout", echoReadTime)
+ mux.HandleFunc("/global", global)
  mux.HandleFunc("/no-context-takeover-decompression", echoNoContextDecompression)
  mux.HandleFunc("/no-context-takeover-decompression-and-compression", echoNoContextDecompressionAndCompression)
  mux.HandleFunc("/context-takeover-decompression", echoContextTakeoverDecompression)

diff --git a/client.go b/client.go
@@ -287,7 +287,10 @@ func (d *DialOption) Dial() (wsCon *Conn, err error) {
  if err := conn.SetDeadline(time.Time{}); err != nil {
  return nil, err
  }
- wsCon = newConn(conn, true /* client is true*/, &d.Config, fr, br)
+ if wsCon, err = newConn(conn, true /* client is true*/, &d.Config, fr, br); err != nil {
+ return nil, err
+ }
  wsCon.pd = pd
+ wsCon.Callback = d.cb
  return wsCon, nil
 }
diff --git a/common_options.go b/common_options.go
@@ -23,7 +23,7 @@ import (
 // 0. CallbackFunc
 func WithClientCallbackFunc(open OnOpenFunc, m OnMessageFunc, c OnCloseFunc) ClientOption {
  return func(o *DialOption) {
- o.Callback = &funcToCallback{
+ o.cb = &funcToCallback{
  onOpen: open,
  onMessage: m,
  onClose: c,
@@ -34,7 +34,7 @@ func WithClientCallbackFunc(open OnOpenFunc, m OnMessageFunc, c OnCloseFunc) Cli
 // 配置服务端回调函数
 func WithServerCallbackFunc(open OnOpenFunc, m OnMessageFunc, c OnCloseFunc) ServerOption {
  return func(o *ConnOption) {
- o.Callback = &funcToCallback{
+ o.cb = &funcToCallback{
  onOpen: open,
  onMessage: m,
  onClose: c,
@@ -46,14 +46,14 @@ func WithServerCallbackFunc(open OnOpenFunc, m OnMessageFunc, c OnCloseFunc) Ser
 // 配置客户端callback
 func WithClientCallback(cb Callback) ClientOption {
  return func(o *DialOption) {
- o.Callback = cb
+ o.cb = cb
  }
 }
 
 // 配置服务端回调函数
 func WithServerCallback(cb Callback) ServerOption {
  return func(o *ConnOption) {
- o.Callback = cb
+ o.cb = cb
  }
 }
 
@@ -90,14 +90,14 @@ func WithClientEnableUTF8Check() ClientOption {
 // 仅仅配置OnMessae函数
 func WithServerOnMessageFunc(cb OnMessageFunc) ServerOption {
  return func(o *ConnOption) {
- o.Callback = OnMessageFunc(cb)
+ o.cb = OnMessageFunc(cb)
  }
 }
 
 // 仅仅配置OnMessae函数
 func WithClientOnMessageFunc(cb OnMessageFunc) ClientOption {
  return func(o *DialOption) {
- o.Callback = OnMessageFunc(cb)
+ o.cb = OnMessageFunc(cb)
  }
 }
 
@@ -292,14 +292,14 @@ func WithClientReadTimeout(t time.Duration) ClientOption {
 // 17.1 配置服务端OnClose
 func WithServerOnCloseFunc(onClose func(c *Conn, err error)) ServerOption {
  return func(o *ConnOption) {
- o.Callback = OnCloseFunc(onClose)
+ o.cb = OnCloseFunc(onClose)
  }
 }
 
 // 17.2 配置客户端OnClose
 func WithClientOnCloseFunc(onClose func(c *Conn, err error)) ClientOption {
  return func(o *DialOption) {
- o.Callback = OnCloseFunc(onClose)
+ o.cb = OnCloseFunc(onClose)
  }
 }
 

diff --git a/config.go b/config.go
@@ -41,7 +41,7 @@ type DialerTimeout interface {
 // 一种是声明一个全局的配置，后面不停使用。
 // 另外一种是局部声明一个配置，然后使用WithXXX函数设置配置
 type Config struct {
- Callback
+ cb Callback
  deflate.PermessageDeflateConf // 静态配置, 从WithXXX函数中获取
  tcpNoDelay bool
  replyPing bool // 开启自动回复
@@ -67,7 +67,7 @@ func (c *Config) initPayloadSize() int {
 
 // 默认设置
 func (c *Config) defaultSetting() error {
- c.Callback = &DefCallback{}
+ c.cb = &DefCallback{}
  c.maxDelayWriteNum = 10
  c.windowsMultipleTimesPayloadSize = 1.0
  c.delayWriteInitBufferSize = 8 * 1024

diff --git a/config_test.go b/config_test.go
@@ -67,7 +67,7 @@ func TestConfig_defaultSetting(t *testing.T) {
  for _, tt := range tests {
  t.Run(tt.name, func(t *testing.T) {
  c := &Config{
- Callback: tt.fields.Callback,
+ cb:  tt.fields.Callback,
  tcpNoDelay: tt.fields.tcpNoDelay,
  replyPing: tt.fields.replyPing,
  ignorePong: tt.fields.ignorePong,

diff --git a/conn.go b/conn.go
@@ -58,6 +58,7 @@ type delayWrite struct {
 type Conn struct {
  fr fixedreader.FixedReader // 默认使用windows
  c net.Conn // net.Conn
+ Callback // callback移至conn中
  br *bufio.Reader // read和fr同时只能使用一个
  *Config // config 可能是全局，也可能是局部初始化得来的
  pd deflate.PermessageDeflateConf // permessageDeflate局部配置
@@ -87,18 +88,20 @@ func setNoDelay(c net.Conn, noDelay bool) error {
  return nil
 }
 
-func newConn(c net.Conn, client bool, conf *Config, fr fixedreader.FixedReader, br *bufio.Reader) *Conn {
- _ = setNoDelay(c, conf.tcpNoDelay)
+func newConn(c net.Conn, client bool, conf *Config, fr fixedreader.FixedReader, br *bufio.Reader) (wsCon *Conn, err error) {
+ if err = setNoDelay(c, conf.tcpNoDelay); err != nil {
+ return nil, err
+ }
 
- con := &Conn{
+ wsCon = &Conn{
  c: c,
  client: client,
  Config: conf,
  fr: fr,
  br: br,
  }
 
- return con
+ return wsCon, err
 }
 
 // 返回标准库的net.Conn

diff --git a/upgrade.go b/upgrade.go
@@ -43,7 +43,11 @@ func NewUpgrade(opts ...ServerOption) *UpgradeServer {
 }
 
 func (u *UpgradeServer) Upgrade(w http.ResponseWriter, r *http.Request) (c *Conn, err error) {
- return upgradeInner(w, r, &u.config)
+ return upgradeInner(w, r, &u.config, nil)
+}
+
+func (u *UpgradeServer) UpgradeV2(w http.ResponseWriter, r *http.Request, cb Callback) (c *Conn, err error) {
+ return upgradeInner(w, r, &u.config, cb)
 }
 
 func Upgrade(w http.ResponseWriter, r *http.Request, opts ...ServerOption) (c *Conn, err error) {
@@ -54,10 +58,10 @@ func Upgrade(w http.ResponseWriter, r *http.Request, opts ...ServerOption) (c *C
  for _, o := range opts {
  o(&conf)
  }
- return upgradeInner(w, r, &conf.Config)
+ return upgradeInner(w, r, &conf.Config, nil)
 }
 
-func upgradeInner(w http.ResponseWriter, r *http.Request, conf *Config) (c *Conn, err error) {
+func upgradeInner(w http.ResponseWriter, r *http.Request, conf *Config, cb Callback) (wsCon *Conn, err error) {
  if ecode, err := checkRequest(r); err != nil {
  http.Error(w, err.Error(), ecode)
  return nil, err
@@ -125,9 +129,15 @@ func upgradeInner(w http.ResponseWriter, r *http.Request, conf *Config) (c *Conn
  if err := conn.SetDeadline(time.Time{}); err != nil {
  return nil, err
  }
- wsCon := newConn(conn, false, conf, fr, br)
+ if wsCon, err = newConn(conn, false, conf, fr, br); err != nil {
+ return nil, err
+ }
 
  wsCon.pd = pd
+ wsCon.Callback = cb
+ if cb == nil {
+ wsCon.Callback = conf.cb
+ }
  return wsCon, nil
 }