35 lines
1004 B
Go
35 lines
1004 B
Go
package server
|
|
|
|
// Utilities for server package: WebSocket upgrader with Origin checks.
|
|
import (
|
|
"net/http"
|
|
"strings"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
// Upgrader creates a Gorilla websocket.Upgrader that enforces Origin checks.
|
|
// If bypass is true, any Origin is accepted. Otherwise, the Origin must match
|
|
// the provided originAddress (host[:port]).
|
|
func Upgrader(originAddress string, bypass bool) websocket.Upgrader {
|
|
if strings.HasPrefix(originAddress, "http://") || strings.HasPrefix(originAddress, "https://") {
|
|
originAddress = strings.TrimPrefix(originAddress, "http://")
|
|
originAddress = strings.TrimPrefix(originAddress, "https://")
|
|
}
|
|
return websocket.Upgrader{
|
|
CheckOrigin: func(r *http.Request) bool {
|
|
if bypass {
|
|
return true
|
|
}
|
|
origin := r.Header.Get("Origin")
|
|
if origin == "" {
|
|
return false
|
|
}
|
|
if strings.HasPrefix(origin, "http://"+originAddress) || strings.HasPrefix(origin, "https://"+originAddress) {
|
|
return true
|
|
}
|
|
return false
|
|
},
|
|
}
|
|
}
|