<commit_before>package main

import (
	"code.google.com/p/go.net/websocket"
	"encoding/json"
	"flag"
	"fmt"
	"gopkg.in/redis.v1"
	"log"
	"net/http"
	"net/url"
	"strings"
)

var (
	redisHost                 string
	redisPort                 int
	redisChannel              string
	redisEndpoint             string
	websocketHost             string
	websocketPort             int
	websocketEndpoint         string
	websocketRoute            string
	websocketAllowableOrigins string
	websocketAllowableURLs    []url.URL
	redisClient               *redis.Client
)

func pubsocketdHandler(w http.ResponseWriter, req *http.Request) {

	// config := pubsocketdConfig{Origin: websocketAllowableURLs}
	// log.Printf("config %v", config)

	originURL := websocketAllowableURLs[0]

	pubsocketdConfig := websocket.Config{Origin: &originURL}

	s := websocket.Server{
		Config:    pubsocketdConfig,
		Handshake: pubsocketdHandshake,
		Handler:   websocket.Handler(pubSubHandler),
	}

	s.ServeHTTP(w, req)
}

func pubsocketdHandshake(config *websocket.Config, req *http.Request) (err error) {

	remoteAddr := req.RemoteAddr
	headers := req.Header

	origin := headers.Get("Origin")
	realIP := headers.Get("X-Real-IP")

	if origin == "" {
		log.Printf("[%s][%s][handshake] missing origin", realIP, remoteAddr)
		return fmt.Errorf("missing origin")
	}

	parsed, err := url.Parse(origin)

	if err != nil {
		log.Printf("[%s][%s][handshake] failed to parse origin, %v", realIP, remoteAddr, origin)
		return fmt.Errorf("invalid origin")
	}

	// See above inre: config.Origin being/becoming a list of url.URLs
	// (20140727/straup)

	if parsed.String() != config.Origin.String() {
		log.Printf("[%s][%s][handshake] invalid origin, %v", realIP, remoteAddr, parsed)
		return fmt.Errorf("invalid origin")
	}

	log.Printf("[%s][%s][handshake] OK", realIP, remoteAddr)
	return
}

func pubSubHandler(ws *websocket.Conn) {

	remoteAddr := ws.Request().RemoteAddr
	headers := ws.Request().Header

	realIP := headers.Get("X-Real-IP")

	log.Printf("[%s][%s][request] OK", realIP, remoteAddr)

	pubsubClient := redisClient.PubSub()
	defer pubsubClient.Close()

	if err := pubsubClient.Subscribe(redisChannel); err != nil {
		log.Printf("[%s][%s][error] failed to subscribe to pubsub channel %v, because %s", realIP, remoteAddr, redisChannel, err)
		ws.Close()
		return
	}

	log.Printf("[%s][%s][connect] OK", realIP, remoteAddr)

	for ws != nil {

		i, _ := pubsubClient.Receive()

		if msg, _ := i.(*redis.Message); msg != nil {

			// log.Printf("[%s][%s][send] %s", realIP, remoteAddr, msg.Payload)

			var json_blob interface{}
			bytes_blob := []byte(msg.Payload)

			if err := json.Unmarshal(bytes_blob, &json_blob); err != nil {
				log.Printf("[%s][%s][error] failed to parse JSON %v, because %v", realIP, remoteAddr, msg.Payload, err)
				continue
			}

			if err := websocket.JSON.Send(ws, json_blob); err != nil {
				log.Printf("[%s][%s][error] failed to send JSON, because %v", realIP, remoteAddr, err)
				ws.Close()
				break
			}

			log.Printf("[%s][%s][send] OK", realIP, remoteAddr)
		}
	}
}

func main() {

	flag.StringVar(&websocketHost, "ws-host", "127.0.0.1", "Websocket host")
	flag.IntVar(&websocketPort, "ws-port", 8080, "Websocket port")
	flag.StringVar(&websocketRoute, "ws-route", "/", "Websocket route")
	flag.StringVar(&websocketAllowableOrigins, "ws-origin", "", "Websocket allowable origins")

	flag.StringVar(&redisHost, "rs-host", "127.0.0.1", "Redis host")
	flag.IntVar(&redisPort, "rs-port", 6379, "Redis port")
	flag.StringVar(&redisChannel, "rs-channel", "pubsocketd", "Redis channel")

	flag.Parse()

	if websocketAllowableOrigins == "" {
		err, _ := fmt.Printf("Missing allowable Origin (-ws-origin=http://example.com)")
		panic(err)
	}

	allowed := strings.Split(websocketAllowableOrigins, ",")
	count := len(allowed)

	if count > 1 {
		err, _ := fmt.Printf("Only one origin server is supported at the moment")
		panic(err)
	}

	websocketAllowableURLs = make([]url.URL, 0, count)

	for _, test := range allowed {

		test := strings.TrimSpace(test)

		url, err := url.Parse(test)

		if err != nil {
			err, _ := fmt.Printf("Invalid Origin parameter: %v, %v", test, err)
			panic(err)
		}

		websocketAllowableURLs = append(websocketAllowableURLs, *url)
	}

	websocketEndpoint = fmt.Sprintf("%s:%d", websocketHost, websocketPort)
	redisEndpoint = fmt.Sprintf("%s:%d", redisHost, redisPort)

	redisClient = redis.NewTCPClient(&redis.Options{
		Addr: redisEndpoint,
	})

	defer redisClient.Close()

	// Normally this is the sort of thing you'd expect to do:
	// http.Handle(websocketRoute, websocket.Handler(pubSubHandler))

	// However since we're going to be aggressively paranoid about checking
	// the Origin headers we're going to set up our own websocket Server
	// thingy complete with custom Config and Handshake directive and
	// pass the whole thing off to HandleFunc (20140727/straup)

	// See also:
	// http://www.christian-schneider.net/CrossSiteWebSocketHijacking.html
	// https://code.google.com/p/go/source/browse/websocket/server.go?repo=net

	http.HandleFunc(websocketRoute, pubsocketdHandler)

	log.Printf("[init] listening for websocket requests on " + websocketEndpoint + websocketRoute + ", from " + websocketAllowableOrigins)
	log.Printf("[init] listening for pubsub messages from " + redisEndpoint + " sent to the " + redisChannel + " channel")

	if err := http.ListenAndServe(websocketEndpoint, nil); err != nil {
		err, _ := fmt.Printf("Failed to start websocket server, because %v", err)
		panic(err)
	}
}
<commit_msg>notes and comments; still only a single Origin URL...<commit_after>package main

import (
	"code.google.com/p/go.net/websocket"
	"encoding/json"
	"flag"
	"fmt"
	"gopkg.in/redis.v1"
	"log"
	"net/http"
	"net/url"
	"strings"
)

var (
	redisHost                 string
	redisPort                 int
	redisChannel              string
	redisEndpoint             string
	websocketHost             string
	websocketPort             int
	websocketEndpoint         string
	websocketRoute            string
	websocketAllowableOrigins string
	websocketAllowableURLs    []url.URL
	redisClient               *redis.Client
)

// What I'd really like to do is pass in a list of allowable Origin
// URLs but the signature for websocket.Config only allows a single
// url.URL thingy and trying to subclass it resulted in a cascading
// series of "unexpected yak" style errors from the compilers. It
// seems like something that must be possible but my Go-fu is still
// weak... (20140729/straup)

// type pubsocketdConfig struct {
//     websocket.Config
//     Origin []url.URL
// }

func pubsocketdHandler(w http.ResponseWriter, req *http.Request) {

	// See above (20140729/straup)

	originURL := websocketAllowableURLs[0]

	pubsocketdConfig := websocket.Config{Origin: &originURL}

	s := websocket.Server{
		Config:    pubsocketdConfig,
		Handshake: pubsocketdHandshake,
		Handler:   websocket.Handler(pubSubHandler),
	}

	s.ServeHTTP(w, req)
}

func pubsocketdHandshake(config *websocket.Config, req *http.Request) (err error) {

	remoteAddr := req.RemoteAddr
	headers := req.Header

	origin := headers.Get("Origin")
	realIP := headers.Get("X-Real-IP")

	if origin == "" {
		log.Printf("[%s][%s][handshake] missing origin", realIP, remoteAddr)
		return fmt.Errorf("missing origin")
	}

	parsed, err := url.Parse(origin)

	if err != nil {
		log.Printf("[%s][%s][handshake] failed to parse origin, %v", realIP, remoteAddr, origin)
		return fmt.Errorf("failed to parse origin")
	}

	// See above inre: config.Origin being/becoming a list of url.URLs
	// (20140727/straup)

	if parsed.String() != config.Origin.String() {
		log.Printf("[%s][%s][handshake] invalid origin, expected %v but got %v", realIP, remoteAddr, config.Origin, parsed)
		return fmt.Errorf("invalid origin")
	}

	log.Printf("[%s][%s][handshake] OK", realIP, remoteAddr)
	return
}

func pubSubHandler(ws *websocket.Conn) {

	remoteAddr := ws.Request().RemoteAddr
	headers := ws.Request().Header

	realIP := headers.Get("X-Real-IP")

	log.Printf("[%s][%s][request] OK", realIP, remoteAddr)

	pubsubClient := redisClient.PubSub()
	defer pubsubClient.Close()

	if err := pubsubClient.Subscribe(redisChannel); err != nil {
		log.Printf("[%s][%s][error] failed to subscribe to pubsub channel %v, because %s", realIP, remoteAddr, redisChannel, err)
		ws.Close()
		return
	}

	log.Printf("[%s][%s][connect] OK", realIP, remoteAddr)

	for ws != nil {

		i, _ := pubsubClient.Receive()

		if msg, _ := i.(*redis.Message); msg != nil {

			// log.Printf("[%s][%s][send] %s", realIP, remoteAddr, msg.Payload)

			var json_blob interface{}
			bytes_blob := []byte(msg.Payload)

			if err := json.Unmarshal(bytes_blob, &json_blob); err != nil {
				log.Printf("[%s][%s][error] failed to parse JSON %v, because %v", realIP, remoteAddr, msg.Payload, err)
				continue
			}

			if err := websocket.JSON.Send(ws, json_blob); err != nil {
				log.Printf("[%s][%s][error] failed to send JSON, because %v", realIP, remoteAddr, err)
				ws.Close()
				break
			}

			log.Printf("[%s][%s][send] OK", realIP, remoteAddr)
		}
	}
}

func main() {

	flag.StringVar(&websocketHost, "ws-host", "127.0.0.1", "Websocket host")
	flag.IntVar(&websocketPort, "ws-port", 8080, "Websocket port")
	flag.StringVar(&websocketRoute, "ws-route", "/", "Websocket route")
	flag.StringVar(&websocketAllowableOrigins, "ws-origin", "", "Websocket allowable origin")

	flag.StringVar(&redisHost, "rs-host", "127.0.0.1", "Redis host")
	flag.IntVar(&redisPort, "rs-port", 6379, "Redis port")
	flag.StringVar(&redisChannel, "rs-channel", "pubsocketd", "Redis channel")

	flag.Parse()

	if websocketAllowableOrigins == "" {
		err, _ := fmt.Printf("Missing allowable Origin (-ws-origin=http://example.com)")
		panic(err)
	}

	allowed := strings.Split(websocketAllowableOrigins, ",")
	count := len(allowed)

	if count > 1 {
		err, _ := fmt.Printf("Only one origin server is supported at the moment")
		panic(err)
	}

	websocketAllowableURLs = make([]url.URL, 0, count)

	for _, test := range allowed {

		test := strings.TrimSpace(test)

		url, err := url.Parse(test)

		if err != nil {
			err, _ := fmt.Printf("Invalid Origin parameter: %v, %v", test, err)
			panic(err)
		}

		websocketAllowableURLs = append(websocketAllowableURLs, *url)
	}

	websocketEndpoint = fmt.Sprintf("%s:%d", websocketHost, websocketPort)
	redisEndpoint = fmt.Sprintf("%s:%d", redisHost, redisPort)

	redisClient = redis.NewTCPClient(&redis.Options{
		Addr: redisEndpoint,
	})

	defer redisClient.Close()

	// Normally this is the sort of thing you'd expect to do:
	// http.Handle(websocketRoute, websocket.Handler(pubSubHandler))

	// However since we're going to be aggressively paranoid about checking
	// the Origin headers we're going to set up our own websocket Server
	// thingy complete with custom Config and Handshake directive and
	// pass the whole thing off to HandleFunc (20140727/straup)

	// See also:
	// http://www.christian-schneider.net/CrossSiteWebSocketHijacking.html
	// https://code.google.com/p/go/source/browse/websocket/server.go?repo=net

	http.HandleFunc(websocketRoute, pubsocketdHandler)

	log.Printf("[init] listening for websocket requests on " + websocketEndpoint + websocketRoute + ", from " + websocketAllowableOrigins)
	log.Printf("[init] listening for pubsub messages from " + redisEndpoint + " sent to the " + redisChannel + " channel")

	if err := http.ListenAndServe(websocketEndpoint, nil); err != nil {
		err, _ := fmt.Printf("Failed to start websocket server, because %v", err)
		panic(err)
	}
}

{'fraction_non_alphanumeric': 0.14705628079882424, 'fraction_numerical': 0.008299472637676148, 'mean_word_length': 3.975483870967742, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 4, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2916549', 'n_tokens_mistral': 4218, 'n_tokens_neox': 3928, 'n_words': 949}