github.com/JFJun/bsc@v1.0.0/rpc/websocket.go (about) 1 // Copyright 2015 The go-ethereum Authors 2 // This file is part of the go-ethereum library. 3 // 4 // The go-ethereum library is free software: you can redistribute it and/or modify 5 // it under the terms of the GNU Lesser General Public License as published by 6 // the Free Software Foundation, either version 3 of the License, or 7 // (at your option) any later version. 8 // 9 // The go-ethereum library is distributed in the hope that it will be useful, 10 // but WITHOUT ANY WARRANTY; without even the implied warranty of 11 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 // GNU Lesser General Public License for more details. 13 // 14 // You should have received a copy of the GNU Lesser General Public License 15 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>. 16 17 package rpc 18 19 import ( 20 "context" 21 "encoding/base64" 22 "fmt" 23 "net/http" 24 "net/url" 25 "os" 26 "strings" 27 "sync" 28 29 "github.com/JFJun/bsc/log" 30 mapset "github.com/deckarep/golang-set" 31 "github.com/gorilla/websocket" 32 ) 33 34 const ( 35 wsReadBuffer = 1024 36 wsWriteBuffer = 1024 37 ) 38 39 var wsBufferPool = new(sync.Pool) 40 41 // WebsocketHandler returns a handler that serves JSON-RPC to WebSocket connections. 42 // 43 // allowedOrigins should be a comma-separated list of allowed origin URLs. 44 // To allow connections with any origin, pass "*". 45 func (s *Server) WebsocketHandler(allowedOrigins []string) http.Handler { 46 var upgrader = websocket.Upgrader{ 47 ReadBufferSize: wsReadBuffer, 48 WriteBufferSize: wsWriteBuffer, 49 WriteBufferPool: wsBufferPool, 50 CheckOrigin: wsHandshakeValidator(allowedOrigins), 51 } 52 return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 53 conn, err := upgrader.Upgrade(w, r, nil) 54 if err != nil { 55 log.Debug("WebSocket upgrade failed", "err", err) 56 return 57 } 58 codec := newWebsocketCodec(conn) 59 s.ServeCodec(codec, 0) 60 }) 61 } 62 63 // wsHandshakeValidator returns a handler that verifies the origin during the 64 // websocket upgrade process. When a '*' is specified as an allowed origins all 65 // connections are accepted. 66 func wsHandshakeValidator(allowedOrigins []string) func(*http.Request) bool { 67 origins := mapset.NewSet() 68 allowAllOrigins := false 69 70 for _, origin := range allowedOrigins { 71 if origin == "*" { 72 allowAllOrigins = true 73 } 74 if origin != "" { 75 origins.Add(strings.ToLower(origin)) 76 } 77 } 78 // allow localhost if no allowedOrigins are specified. 79 if len(origins.ToSlice()) == 0 { 80 origins.Add("http://localhost") 81 if hostname, err := os.Hostname(); err == nil { 82 origins.Add("http://" + strings.ToLower(hostname)) 83 } 84 } 85 log.Debug(fmt.Sprintf("Allowed origin(s) for WS RPC interface %v", origins.ToSlice())) 86 87 f := func(req *http.Request) bool { 88 // Skip origin verification if no Origin header is present. The origin check 89 // is supposed to protect against browser based attacks. Browsers always set 90 // Origin. Non-browser software can put anything in origin and checking it doesn't 91 // provide additional security. 92 if _, ok := req.Header["Origin"]; !ok { 93 return true 94 } 95 // Verify origin against whitelist. 96 origin := strings.ToLower(req.Header.Get("Origin")) 97 if allowAllOrigins || origins.Contains(origin) { 98 return true 99 } 100 log.Warn("Rejected WebSocket connection", "origin", origin) 101 return false 102 } 103 104 return f 105 } 106 107 type wsHandshakeError struct { 108 err error 109 status string 110 } 111 112 func (e wsHandshakeError) Error() string { 113 s := e.err.Error() 114 if e.status != "" { 115 s += " (HTTP status " + e.status + ")" 116 } 117 return s 118 } 119 120 // DialWebsocketWithDialer creates a new RPC client that communicates with a JSON-RPC server 121 // that is listening on the given endpoint using the provided dialer. 122 func DialWebsocketWithDialer(ctx context.Context, endpoint, origin string, dialer websocket.Dialer) (*Client, error) { 123 endpoint, header, err := wsClientHeaders(endpoint, origin) 124 if err != nil { 125 return nil, err 126 } 127 return newClient(ctx, func(ctx context.Context) (ServerCodec, error) { 128 conn, resp, err := dialer.DialContext(ctx, endpoint, header) 129 if err != nil { 130 hErr := wsHandshakeError{err: err} 131 if resp != nil { 132 hErr.status = resp.Status 133 } 134 return nil, hErr 135 } 136 return newWebsocketCodec(conn), nil 137 }) 138 } 139 140 // DialWebsocket creates a new RPC client that communicates with a JSON-RPC server 141 // that is listening on the given endpoint. 142 // 143 // The context is used for the initial connection establishment. It does not 144 // affect subsequent interactions with the client. 145 func DialWebsocket(ctx context.Context, endpoint, origin string) (*Client, error) { 146 dialer := websocket.Dialer{ 147 ReadBufferSize: wsReadBuffer, 148 WriteBufferSize: wsWriteBuffer, 149 WriteBufferPool: wsBufferPool, 150 } 151 return DialWebsocketWithDialer(ctx, endpoint, origin, dialer) 152 } 153 154 func wsClientHeaders(endpoint, origin string) (string, http.Header, error) { 155 endpointURL, err := url.Parse(endpoint) 156 if err != nil { 157 return endpoint, nil, err 158 } 159 header := make(http.Header) 160 if origin != "" { 161 header.Add("origin", origin) 162 } 163 if endpointURL.User != nil { 164 b64auth := base64.StdEncoding.EncodeToString([]byte(endpointURL.User.String())) 165 header.Add("authorization", "Basic "+b64auth) 166 endpointURL.User = nil 167 } 168 return endpointURL.String(), header, nil 169 } 170 171 func newWebsocketCodec(conn *websocket.Conn) ServerCodec { 172 conn.SetReadLimit(maxRequestContentLength) 173 return NewFuncCodec(conn, conn.WriteJSON, conn.ReadJSON) 174 }