photoview/api/server/websocket.go

package server

import (
	"log"
	"net/http"
	"net/url"

	"github.com/gorilla/websocket"
	"github.com/photoview/photoview/api/utils"
)

func WebsocketUpgrader(devMode bool) websocket.Upgrader {
	return websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			if devMode {
				return true
			} else {
				uiEndpoint := utils.UiEndpointUrl()
				if uiEndpoint == nil {
					return true
				}

				if r.Header.Get("origin") == "" {
					return true
				}

				originURL, err := url.Parse(r.Header.Get("origin"))
				if err != nil {
					log.Printf("Could not parse origin header of websocket request: %s", err)
					return false
				}

				return isUIOnSameHost(uiEndpoint, originURL)
			}
		},
	}
}

func isUIOnSameHost(uiEndpoint *url.URL, originURL *url.URL) bool {
	if uiEndpoint.Host == originURL.Host {
		return true
	} else {
		log.Printf("Not allowing websocket request from %s because it doesn't match PHOTOVIEW_UI_ENDPOINT %s",
			originURL.Host, uiEndpoint.Host)
		return false
	}
}