package client import ( "bufio" "crypto/rand" "crypto/sha256" "encoding/base64" "encoding/binary" "fmt" "io" "net" "net/url" "strings" "time" ) const ( wsGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11" wsVersion = "13" crlf = "\r\n" ) // wsConn 封装一条 WebSocket 连接(客户端视角,帧带 mask)。 type wsConn struct { conn net.Conn br *bufio.Reader bw *bufio.Writer closed bool } // dialWS 发起 WS 客户端握手升级。 // 支持 ws:// 和 wss://(wss 暂未实现,若需要需加 TLS dial)。 func dialWS(rawURL, token string, timeout time.Duration) (*wsConn, error) { u, err := url.Parse(rawURL) if err != nil { return nil, fmt.Errorf("devicebridge: invalid ws url %q: %w", rawURL, err) } host := u.Host if u.Port() == "" { if u.Scheme == "wss" { host = host + ":443" } else { host = host + ":80" } } path := u.Path if u.RawQuery != "" { path = path + "?" + u.RawQuery } if path == "" { path = "/" } // 默认路径 if token != "" && strings.Index(path, "token=") < 0 { if strings.ContainsRune(path, '?') { path = path + "&token=" + urlEscape(token) } else { path = path + "?token=" + urlEscape(token) } } dialer := net.Dialer{Timeout: timeout} conn, err := dialer.Dial("tcp", host) if err != nil { return nil, fmt.Errorf("devicebridge: dial %s: %w", host, err) } key := wsKey() var sb strings.Builder sb.WriteString("GET " + path + " HTTP/1.1" + crlf) sb.WriteString("Host: " + host + crlf) sb.WriteString("Upgrade: websocket" + crlf) sb.WriteString("Connection: Upgrade" + crlf) sb.WriteString("Sec-WebSocket-Key: " + key + crlf) sb.WriteString("Sec-WebSocket-Version: " + wsVersion + crlf + crlf) if _, err := conn.Write([]byte(sb.String())); err != nil { conn.Close() return nil, fmt.Errorf("devicebridge: write upgrade: %w", err) } br := bufio.NewReader(conn) var headerBuf strings.Builder for { line, err := br.ReadString('\n') if err != nil { conn.Close() return nil, fmt.Errorf("devicebridge: read upgrade resp: %w", err) } headerBuf.WriteString(line) if strings.Contains(headerBuf.String(), crlf+crlf) { break } } if !strings.Contains(headerBuf.String(), " 101 ") { conn.Close() return nil, fmt.Errorf("devicebridge: upgrade failed: %s", firstLine(headerBuf.String())) } return &wsConn{conn: conn, br: br, bw: bufio.NewWriter(conn)}, nil } // writeText 发送 WS 文本帧(0x1,带 mask)。 func (w *wsConn) writeText(payload []byte) error { return w.writeFrame(0x1, payload) } // writeBinary 发送 WS 二进制帧(0x2,带 mask)。 func (w *wsConn) writeBinary(payload []byte) error { return w.writeFrame(0x2, payload) } // writeFrame 发送一个 WS 帧(客户端 mask 模式)。 func (w *wsConn) writeFrame(opcode byte, payload []byte) error { if w.closed { return fmt.Errorf("devicebridge: connection closed") } length := len(payload) // 帧头 hdrLen := 2 switch { case length < 126: // 1 byte length case length <= 0xffff: hdrLen += 2 default: hdrLen += 8 } hdrLen += 4 // mask key hdr := make([]byte, hdrLen) hdr[0] = 0x80 | opcode switch { case length < 126: hdr[1] = 0x80 | byte(length) case length <= 0xffff: hdr[1] = 0x80 | 126 binary.BigEndian.PutUint16(hdr[2:4], uint16(length)) default: hdr[1] = 0x80 | 127 binary.BigEndian.PutUint64(hdr[2:10], uint64(length)) } // mask key var maskKey [4]byte rand.Read(maskKey[:]) copy(hdr[hdrLen-4:], maskKey[:]) // mask payload masked := make([]byte, length) for i := 0; i < length; i++ { masked[i] = payload[i] ^ maskKey[i&3] } if _, err := w.bw.Write(hdr); err != nil { return err } if _, err := w.bw.Write(masked); err != nil { return err } return w.bw.Flush() } // readFrame 读取一个 WS 帧。返回 (payload, isClose, opcode, error)。 // 客户端收到的帧不带 mask。 func (w *wsConn) readFrame() ([]byte, bool, byte, error) { if w.closed { return nil, true, 0, fmt.Errorf("devicebridge: connection closed") } b0, err := w.br.ReadByte() if err != nil { return nil, true, 0, err } opcode := b0 & 0x0f b1, err := w.br.ReadByte() if err != nil { return nil, true, 0, err } length := uint64(b1 & 0x7f) if length == 126 { var ext [2]byte if _, err := io.ReadFull(w.br, ext[:]); err != nil { return nil, true, opcode, err } length = uint64(binary.BigEndian.Uint16(ext[:])) } else if length == 127 { var ext [8]byte if _, err := io.ReadFull(w.br, ext[:]); err != nil { return nil, true, opcode, err } length = binary.BigEndian.Uint64(ext[:]) } // 二进制帧允许更大(8MB),文本帧 1MB maxFrame := uint64(1 << 20) if opcode == 0x2 { maxFrame = 8 << 20 } if length > maxFrame { return nil, true, opcode, fmt.Errorf("devicebridge: frame too large (%d bytes)", length) } payload := make([]byte, length) if _, err := io.ReadFull(w.br, payload); err != nil { return nil, true, opcode, err } switch opcode { case 0x1, 0x2: return payload, false, opcode, nil case 0x8: return nil, true, opcode, nil case 0x9: // ping return nil, false, opcode, errPing case 0xa: // pong return nil, false, opcode, nil default: return nil, false, opcode, fmt.Errorf("devicebridge: unsupported opcode %x", opcode) } } // writePong 发送 pong 帧。 func (w *wsConn) writePong() error { return w.writeFrame(0xa, nil) } // close 发送关闭帧并关闭连接。 func (w *wsConn) close() error { w.closed = true _ = w.writeFrame(0x8, nil) return w.conn.Close() } // ===== 辅助函数 ===== var errPing = fmt.Errorf("ping") func wsKey() string { var b [16]byte rand.Read(b[:]) return base64.StdEncoding.EncodeToString(b[:]) } func urlEscape(s string) string { var sb strings.Builder const hex = "0123456789ABCDEF" for i := 0; i < len(s); i++ { c := s[i] if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '-' || c == '_' || c == '.' || c == '~' { sb.WriteByte(c) } else { sb.WriteByte('%') sb.WriteByte(hex[c>>4]) sb.WriteByte(hex[c&0xf]) } } return sb.String() } func firstLine(s string) string { if i := strings.IndexByte(s, '\n'); i >= 0 { return strings.TrimSpace(s[:i]) } return strings.TrimSpace(s) } func wsAccept(key string) string { h := sha256.Sum256([]byte(key + wsGUID)) return base64.StdEncoding.EncodeToString(h[:]) } // readWSFrame 读取一个 WS 帧(从已有的 bufio.Reader,兼容非 wsConn 场景)。 func readWSFrame(br *bufio.Reader) ([]byte, bool, byte, error) { b0, err := br.ReadByte() if err != nil { return nil, true, 0, err } opcode := b0 & 0x0f b1, err := br.ReadByte() if err != nil { return nil, true, 0, err } length := uint64(b1 & 0x7f) if length == 126 { var ext [2]byte if _, err := io.ReadFull(br, ext[:]); err != nil { return nil, true, opcode, err } length = uint64(binary.BigEndian.Uint16(ext[:])) } else if length == 127 { var ext [8]byte if _, err := io.ReadFull(br, ext[:]); err != nil { return nil, true, opcode, err } length = binary.BigEndian.Uint64(ext[:]) } maxFrame := uint64(1 << 20) if opcode == 0x2 { maxFrame = 8 << 20 } if length > maxFrame { return nil, true, opcode, fmt.Errorf("frame too large") } payload := make([]byte, length) if _, err := io.ReadFull(br, payload); err != nil { return nil, true, opcode, err } switch opcode { case 0x1, 0x2: return payload, false, opcode, nil case 0x8: return nil, true, opcode, nil default: return nil, false, opcode, nil } } // writeWSFrame 发送一个 WS 帧(非 mask 模式,服务端用)。 func writeWSFrame(w io.Writer, opcode byte, payload []byte) error { length := len(payload) hdr := []byte{0x80 | opcode} switch { case length < 126: hdr = append(hdr, byte(length)) case length <= 0xffff: hdr = append(hdr, 126, 0, 0) binary.BigEndian.PutUint16(hdr[len(hdr)-2:], uint16(length)) default: hdr = append(hdr, 127, 0, 0, 0, 0, 0, 0, 0, 0) binary.BigEndian.PutUint64(hdr[len(hdr)-8:], uint64(length)) } if _, err := w.Write(hdr); err != nil { return err } if _, err := w.Write(payload); err != nil { return err } return nil } // ensureTimeout 设置连接读写超时。 func (w *wsConn) setDeadline(t time.Time) { if w.conn != nil { w.conn.SetDeadline(t) } } // LocalAddr 返回本地地址。 func (w *wsConn) LocalAddr() net.Addr { if w.conn != nil { return w.conn.LocalAddr() } return nil } // RemoteAddr 返回远程地址。 func (w *wsConn) RemoteAddr() net.Addr { if w.conn != nil { return w.conn.RemoteAddr() } return nil }