ws/internal/http/upgrade/header_check.go

package upgrade

import (
	"fmt"
	"strconv"
	"strings"
)

// checkHost checks and extracts the Host header
func (r *Request) extractHostPort(bb HeaderValue) error {

	if len(bb) != 1 {
		return &InvalidRequest{"Host", fmt.Sprintf("expected single value, got %d", len(bb))}
	}

	if len(bb[0]) <= 3 {
		return &InvalidRequest{"Host", fmt.Sprintf("expected non-empty value (got %d bytes)", len(bb[0]))}
	}

	split := strings.Split(string(bb[0]), ":")

	r.host = split[0]

	// no port
	if len(split) < 2 {
		return nil
	}

	// extract port
	readPort, err := strconv.ParseUint(split[1], 10, 16)
	if err != nil {
		r.code = BadRequest
		return &InvalidRequest{"Host", "cannot read port"}
	}

	r.port = uint16(readPort)

	// if 'Origin' header is already read, check it
	if len(r.origin) > 0 {
		if err != nil {
			err = r.checkOriginPolicy()
			r.code = Forbidden
			return &InvalidOriginPolicy{r.host, r.origin, err}
		}
	}

	return nil

}

// checkOrigin checks the Origin Header
func (r *Request) extractOrigin(bb HeaderValue) error {

	// bypass
	if bypassOriginPolicy {
		return nil
	}

	if len(bb) != 1 {
		r.code = Forbidden
		return &InvalidRequest{"Origin", fmt.Sprintf("expected single value, got %d", len(bb))}
	}

	r.origin = string(bb[0])

	// if host already stored, check origin policy
	if len(r.host) > 0 {
		err := r.checkOriginPolicy()
		if err != nil {
			r.code = Forbidden
			return &InvalidOriginPolicy{r.host, r.origin, err}
		}
	}

	return nil

}

// checkOriginPolicy origin policy based on 'host' value
func (r *Request) checkOriginPolicy() error {
	// TODO: Origin policy, for now BYPASS
	r.validPolicy = true
	return nil
}

// checkConnection checks the 'Connection' header
// it MUST contain 'Upgrade'
func (r *Request) checkConnection(bb HeaderValue) error {

	for _, b := range bb {

		if strings.ToLower(string(b)) == "upgrade" {
			r.hasConnection = true
			return nil
		}

	}

	r.code = BadRequest
	return &InvalidRequest{"Upgrade", "expected 'Upgrade' (case insensitive)"}

}

// checkUpgrade checks the 'Upgrade' header
// it MUST be 'websocket'
func (r *Request) checkUpgrade(bb HeaderValue) error {

	if len(bb) != 1 {
		r.code = BadRequest
		return &InvalidRequest{"Upgrade", fmt.Sprintf("expected single value, got %d", len(bb))}
	}

	if strings.ToLower(string(bb[0])) == "websocket" {
		r.hasUpgrade = true
		return nil
	}

	r.code = BadRequest
	return &InvalidRequest{"Upgrade", fmt.Sprintf("expected 'websocket' (case insensitive), got '%s'", bb[0])}

}

// checkVersion checks the 'Sec-WebSocket-Version' header
// it MUST be '13'
func (r *Request) checkVersion(bb HeaderValue) error {

	if len(bb) != 1 || string(bb[0]) != "13" {
		r.code = UpgradeRequired
		return &InvalidRequest{"Sec-WebSocket-Version", fmt.Sprintf("expected '13', got '%s'", bb[0])}
	}

	r.hasVersion = true
	return nil

}

// extractKey extracts the 'Sec-WebSocket-Key' header
// it MUST be 24 bytes (base64)
func (r *Request) extractKey(bb HeaderValue) error {

	if len(bb) != 1 || len(bb[0]) != 24 {
		r.code = BadRequest
		return &InvalidRequest{"Sec-WebSocket-Key", fmt.Sprintf("expected 24 bytes base64, got %d bytes", len(bb[0]))}
	}

	r.key = bb[0]

	return nil

}

// extractProtocols extracts the 'Sec-WebSocket-Protocol' header
// it can contain multiple values
func (r *Request) extractProtocols(bb HeaderValue) error {

	r.protocols = bb

	return nil
}