ws/internal/http/upgrade/response.go

package upgrade

import (
	"crypto/sha1"
	"encoding/base64"
	"fmt"
	"io"
)

// constants
const (
	httpVersion = "1.1"
	wsVersion   = 13
	keySalt     = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
)

// Response is an HTTP Upgrade Response
type Response struct {
	StatusCode StatusCode
	// Sec-WebSocket-Protocol or nil if missing
	Protocol []byte
	// processed from Sec-WebSocket-Key
	key []byte
}

// ProcessKey processes the accept token according
// to the rfc from the Sec-WebSocket-Key
func (r *Response) ProcessKey(k []byte) {
	// ignore empty key
	if k == nil || len(k) < 1 {
		return
	}

	// concat with constant salt
	salted := append(k, []byte(keySalt)...)
	// hash with sha1
	digest := sha1.Sum(salted)
	// base64 encode
	r.key = []byte(base64.StdEncoding.EncodeToString(digest[:sha1.Size]))
}

// WriteTo writes the response; typically in a socket
//
// implements io.WriterTo
func (r Response) WriteTo(w io.Writer) (int64, error) {
	responseLine := fmt.Sprintf("HTTP/%s %d %s\r\n", httpVersion, r.StatusCode, r.StatusCode)

	optionalProtocol := ""
	if len(r.Protocol) > 0 {
		optionalProtocol = fmt.Sprintf("Sec-WebSocket-Protocol: %s\r\n", r.Protocol)
	}

	headers := fmt.Sprintf("Upgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Version: %d\r\n%s", wsVersion, optionalProtocol)
	if r.key != nil {
		headers = fmt.Sprintf("%sSec-WebSocket-Accept: %s\r\n", headers, r.key)
	}
	headers = fmt.Sprintf("%s\r\n", headers)

	combined := []byte(fmt.Sprintf("%s%s", responseLine, headers))

	written, err := w.Write(combined)
	return int64(written), err

}