fix origin checking

This commit is contained in:
Radon 2025-08-29 14:51:51 -05:00
parent 1c361c218b
commit e199d421f7
2 changed files with 8 additions and 4 deletions

View File

@ -14,7 +14,7 @@ func main() {
var bufSize = flag.Int("bufsize", 256, "Channel buffer size")
var gzipEnabled = flag.Bool("gzip-enable", false, "Enable gzip compression")
var cachingDisabled = flag.Bool("cache-disable", false, "Disable caching")
var originCheck = flag.Bool("origin-check", false, "Enable origin check")
var origin = flag.String("origin", "", "Origin to allow (e.g. example.com), leave blank to allow all")
var help = flag.Bool("help", false, "Show help")
flag.Parse()
@ -26,7 +26,7 @@ func main() {
address := fmt.Sprintf("%s:%d", *ip, *port)
hub := server.NewHub(*bufSize, server.Upgrader(address, !*originCheck))
hub := server.NewHub(*bufSize, server.Upgrader(*origin, *origin == ""))
go Run(hub)
mux := http.NewServeMux()

View File

@ -7,7 +7,11 @@ import (
"github.com/gorilla/websocket"
)
func Upgrader(address string, bypass bool) websocket.Upgrader {
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 {
@ -17,7 +21,7 @@ func Upgrader(address string, bypass bool) websocket.Upgrader {
if origin == "" {
return false
}
if strings.HasPrefix(origin, "http://"+address) || strings.HasPrefix(origin, "https://"+address) {
if strings.HasPrefix(origin, "http://"+originAddress) || strings.HasPrefix(origin, "https://"+originAddress) {
return true
}
return false