Skip to content

Commit

Permalink
feat(ws_reverseproxy): add OptionWS for ws_reverseproxy
Browse files Browse the repository at this point in the history
  • Loading branch information
yeqown committed May 31, 2021
1 parent d442dbe commit 4a03f2f
Show file tree
Hide file tree
Showing 8 changed files with 345 additions and 115 deletions.
1 change: 1 addition & 0 deletions benchmark/ws_proxy_benchmark.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (
"syscall"

"github.com/valyala/fasthttp"

proxy "github.com/yeqown/fasthttp-reverse-proxy/v2"
)

Expand Down
2 changes: 0 additions & 2 deletions debug.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,10 @@ import (
)

var (
debug = true
logger, _ = log.NewLogger()
)

// SetProduction .
func SetProduction() {
debug = false
logger.SetLogLevel(log.LevelInfo)
}
3 changes: 2 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,9 @@ module github.com/yeqown/fasthttp-reverse-proxy/v2

require (
github.com/fasthttp/websocket v1.4.1
github.com/stretchr/testify v1.6.1 // indirect
github.com/valyala/fasthttp v1.4.0
github.com/yeqown/log v1.0.3
github.com/yeqown/log v1.0.5
)

go 1.13
7 changes: 7 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -12,16 +12,23 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/savsgio/gotils v0.0.0-20190714152828-365999d0a274 h1:F52t1X2ziOrMcQMVHo8ZxwOrDTMAq6MrlKtL1Atu2wU=
github.com/savsgio/gotils v0.0.0-20190714152828-365999d0a274/go.mod h1:w803/Fg1m0hrp1ZT9KNfQe4E4+WOMMFLcgzPvOcye10=
github.com/stretchr/objx v0.1.0 h1:4G4v2dO3VZwixGIRoQ5Lfboy6nUhCyYzaqnIAPPhYs4=
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0=
github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=
github.com/valyala/fasthttp v1.4.0 h1:PuaTGZIw3mjYhhhbVbCQp8aciRZN9YdoB7MGX9Ko76A=
github.com/valyala/fasthttp v1.4.0/go.mod h1:4vX61m6KN+xDduDNwXrhIAVZaZaZiQ1luJk8LWSxF3s=
github.com/valyala/tcplisten v0.0.0-20161114210144-ceec8f93295a h1:0R4NLDRDZX6JcmhJgXi5E4b8Wg84ihbmUKp/GvSPEzc=
github.com/valyala/tcplisten v0.0.0-20161114210144-ceec8f93295a/go.mod h1:v3UYOV9WzVtRmSR+PDvWpU/qWl4Wa5LApYYX4ZtKbio=
github.com/yeqown/log v1.0.3 h1:kWwBMytMSkYkr2fmi20PLgDuJNXkq3mOpVbfsTS1soI=
github.com/yeqown/log v1.0.3/go.mod h1:RTslXFTg+8Uj5AizIxdfichvBZi/OOKao6yP3tMtTns=
github.com/yeqown/log v1.0.5 h1:d+zp35OAnJI74bDMdl9QWPssEdPp+SH1WZoy3AhZp3o=
github.com/yeqown/log v1.0.5/go.mod h1:RTslXFTg+8Uj5AizIxdfichvBZi/OOKao6yP3tMtTns=
github.com/yeqown/log v1.0.6-0.20210212131351-ecf09fa100af h1:YixbvpBZW8cdDsAG8rcx84cb/tqx2gk5CcpoqB8Izoc=
github.com/yeqown/log v1.0.6-0.20210212131351-ecf09fa100af/go.mod h1:RTslXFTg+8Uj5AizIxdfichvBZi/OOKao6yP3tMtTns=
golang.org/x/net v0.0.0-20180911220305-26e67e76b6c3 h1:czFLhve3vsQetD6JOJ8NZZvGQIXlnN3/yXxbT6/awxI=
golang.org/x/net v0.0.0-20180911220305-26e67e76b6c3/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
Expand Down
97 changes: 97 additions & 0 deletions ws_reverseporxy_option.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
package proxy

import (
"errors"
"net/http"
"net/url"

"github.com/fasthttp/websocket"
"github.com/valyala/fasthttp"
)

// OptionWS to define all options to reverse web socket proxy.
type OptionWS interface {
apply(o *buildOptionWS)
}

type funcBuildOptionWS struct {
f func(o *buildOptionWS)
}

func newFuncBuildOptionWS(f func(o *buildOptionWS)) funcBuildOptionWS { return funcBuildOptionWS{f: f} }
func (fb funcBuildOptionWS) apply(o *buildOptionWS) { fb.f(o) }

type forwardHeaderHandler func(ctx *fasthttp.RequestCtx) (forwardHeader http.Header)

// buildOptionWS is Option for WS reverse-proxy
type buildOptionWS struct {
// target indicates which backend server to proxy.
target *url.URL

// fn is forwardHeaderHandler which allows users customize themselves' forward headers
// to be proxied to backend server.
fn forwardHeaderHandler

// dialer contains options for connecting to the backend WebSocket server.
// If nil, DefaultDialer is used.
dialer *websocket.Dialer

// upgrader specifies the parameters for upgrading a incoming HTTP
// connection to a WebSocket connection. If nil, DefaultUpgrader is used.
upgrader *websocket.FastHTTPUpgrader
}

func (o *buildOptionWS) validate() error {
if o == nil {
return errors.New("option is nil")
}

if o.target == nil {
return errors.New("target is nil")
}

return nil
}

func defaultBuildOptionWS() *buildOptionWS {
return &buildOptionWS{
target: nil,
fn: nil,
dialer: nil,
upgrader: nil,
}
}

// WithURL_OptionWS specify the url to backend websocket server.
// WithURL_OptionWS("ws://YOUR_WEBSOCKET_HOST:PORT/AND/PATH")
func WithURL_OptionWS(u string) OptionWS {
return newFuncBuildOptionWS(func(o *buildOptionWS) {
URL, err := url.Parse(u)
if err != nil {
panic(err)
}

o.target = URL
})
}

// WithDialer_OptionWS use specified dialer
func WithDialer_OptionWS(dialer *websocket.Dialer) OptionWS {
return newFuncBuildOptionWS(func(o *buildOptionWS) {
o.dialer = dialer
})
}

// WithUpgrader_OptionWS use specified upgrader.
func WithUpgrader_OptionWS(upgrader *websocket.FastHTTPUpgrader) OptionWS {
return newFuncBuildOptionWS(func(o *buildOptionWS) {
o.upgrader = upgrader
})
}

// WithForwardHeadersHandlers_OptionWS allows users to customize forward headers.
func WithForwardHeadersHandlers_OptionWS(handler forwardHeaderHandler) OptionWS {
return newFuncBuildOptionWS(func(o *buildOptionWS) {
o.fn = handler
})
}
42 changes: 42 additions & 0 deletions ws_reverseporxy_option_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
package proxy

import (
"net/http"
"testing"

"github.com/valyala/fasthttp"

"github.com/stretchr/testify/assert"
)

func Test_defaultBuildOptionWS(t *testing.T) {
dst := defaultBuildOptionWS()

assert.Nil(t, dst.target)
assert.Nil(t, dst.dialer)
assert.Nil(t, dst.upgrader)
}

func Test_WithURL_OptionWS(t *testing.T) {
assert.NotPanics(t, func() {
WithURL_OptionWS("ws://localhost:8080/path")
}, "could not panic")

dst := defaultBuildOptionWS()
WithURL_OptionWS("ws://localhost:8080/path").apply(dst)
assert.NotNil(t, dst.target)
assert.Equal(t, "ws", dst.target.Scheme)
assert.Equal(t, "localhost:8080", dst.target.Host)
assert.Equal(t, "8080", dst.target.Port())
assert.Equal(t, "/path", dst.target.Path)
}

func Test_WithForwardHeadersHandlers_OptionWS(t *testing.T) {
dst := defaultBuildOptionWS()
WithForwardHeadersHandlers_OptionWS(func(reqHeader *fasthttp.RequestCtx) (forwardHeader http.Header) {
return http.Header{
"X-Test": []string{"X-Test"},
}
}).apply(dst)
assert.NotNil(t, dst.fn)
}

0 comments on commit 4a03f2f

Please sign in to comment.