ws/internal/http/upgrade/response.go

package upgrade

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

// HTTPVersion constant
const HTTPVersion = "1.1"

// UsedWSVersion constant websocket version
const UsedWSVersion = 13

// WSSalt constant websocket salt
const WSSalt = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

// Response represents an HTTP Upgrade Response
type Response struct {
	code     StatusCode // status code
	accept   []byte     // processed from Sec-WebSocket-Key
	protocol []byte     // set from Sec-WebSocket-Protocol or none if not received
}

// SetStatusCode sets the status code
func (r *Response) SetStatusCode(sc StatusCode) {
	r.code = sc
}

// SetProtocol sets the protocols
func (r *Response) SetProtocol(p []byte) {
	r.protocol = p
}

// ProcessKey processes the accept token according
// to the rfc from the Sec-WebSocket-Key
func (r *Response) ProcessKey(k []byte) {

	// do nothing for empty key
	if k == nil || len(k) == 0 {
		r.accept = nil
		return
	}

	// 1. Concat with constant salt
	mix := append(k, []byte(WSSalt)...)

	// 2. Hash with sha1 algorithm
	digest := sha1.Sum(mix)

	// 3. Base64 encode it
	r.accept = []byte(base64.StdEncoding.EncodeToString(digest[:sha1.Size]))

}

// Send sends the response through an io.Writer
// typically a socket
func (r Response) Send(w io.Writer) (int, error) {

	// 1. Build response line
	responseLine := fmt.Sprintf("HTTP/%s %d %s\r\n", HTTPVersion, r.code, r.code)

	// 2. Build headers
	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", UsedWSVersion, optionalProtocol)
	if r.accept != nil {
		headers = fmt.Sprintf("%sSec-WebSocket-Accept: %s\r\n", headers, r.accept)
	}
	headers = fmt.Sprintf("%s\r\n", headers)

	// 3. Build all
	raw := []byte(fmt.Sprintf("%s%s", responseLine, headers))

	// 4. Write
	written, err := w.Write(raw)

	return written, err

}

// GetProtocol returns the choosen protocol if set, else nil
func (r Response) GetProtocol() []byte {
	return r.protocol
}

// GetStatusCode returns the response status code
func (r Response) GetStatusCode() StatusCode {
	return r.code
}