1 files changed, 75 insertions, 0 deletions

diff --git a/relay/websocket/websocket.go b/relay/websocket/websocket.go
index cfc3289..2ae0dec 100644
--- a/relay/websocket/websocket.go
+++ b/relay/websocket/websocket.go
@@ -7,12 +7,14 @@ import (
         "context"
         "crypto/rand"
         "crypto/sha1"
+        "crypto/tls"
         "encoding/base64"
         "encoding/binary"
         "fmt"
         "io"
         "net"
         "net/http"
+        "net/url"
         "strings"
         "sync"
         "time"
@@ -207,6 +209,79 @@ func acceptKey(key string) string {
         return base64.StdEncoding.EncodeToString(h.Sum(nil))
 }
 
+// Dial connects to a WebSocket server at rawURL and performs the client-side
+// RFC 6455 handshake. Supports ws:// and wss:// schemes.
+func Dial(rawURL string) (*Conn, error) {
+        u, err := url.Parse(rawURL)
+        if err != nil {
+                return nil, fmt.Errorf("websocket: parse url: %w", err)
+        }
+
+        host := u.Host
+        var netConn net.Conn
+        switch u.Scheme {
+        case "ws":
+                if !strings.Contains(host, ":") {
+                        host += ":80"
+                }
+                netConn, err = net.Dial("tcp", host)
+        case "wss":
+                if !strings.Contains(host, ":") {
+                        host += ":443"
+                }
+                netConn, err = tls.Dial("tcp", host, &tls.Config{ServerName: u.Hostname()})
+        default:
+                return nil, fmt.Errorf("websocket: unsupported scheme %q (use ws:// or wss://)", u.Scheme)
+        }
+        if err != nil {
+                return nil, fmt.Errorf("websocket: dial %s: %w", host, err)
+        }
+
+        // Generate a random 16-byte key and base64-encode it.
+        var keyBytes [16]byte
+        if _, err := rand.Read(keyBytes[:]); err != nil {
+                netConn.Close()
+                return nil, fmt.Errorf("websocket: generate key: %w", err)
+        }
+        key := base64.StdEncoding.EncodeToString(keyBytes[:])
+
+        path := u.RequestURI()
+        if path == "" {
+                path = "/"
+        }
+
+        req := "GET " + path + " HTTP/1.1\r\n" +
+                "Host: " + u.Host + "\r\n" +
+                "Upgrade: websocket\r\n" +
+                "Connection: Upgrade\r\n" +
+                "Sec-WebSocket-Key: " + key + "\r\n" +
+                "Sec-WebSocket-Version: 13\r\n\r\n"
+
+        if _, err := netConn.Write([]byte(req)); err != nil {
+                netConn.Close()
+                return nil, fmt.Errorf("websocket: send handshake: %w", err)
+        }
+
+        br := bufio.NewReader(netConn)
+        resp, err := http.ReadResponse(br, nil)
+        if err != nil {
+                netConn.Close()
+                return nil, fmt.Errorf("websocket: read handshake response: %w", err)
+        }
+        resp.Body.Close()
+
+        if resp.StatusCode != 101 {
+                netConn.Close()
+                return nil, fmt.Errorf("websocket: server returned status %d, want 101", resp.StatusCode)
+        }
+        if resp.Header.Get("Sec-WebSocket-Accept") != acceptKey(key) {
+                netConn.Close()
+                return nil, fmt.Errorf("websocket: bad Sec-WebSocket-Accept header")
+        }
+
+        return &Conn{rwc: netConn, br: br, client: true}, nil
+}
+
 // Accept performs the server-side WebSocket handshake, hijacking the HTTP
 // connection and returning a Conn ready for framed I/O.
 func Accept(w http.ResponseWriter, r *http.Request) (*Conn, error) {