teowebrtc_signal.go

// Copyright 2021-2024 Kirill Scherba <kirill@scherba.ru>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Webretc signal server (for teonet network)
package teowebrtc_signal

import (
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
	"github.com/teonet-go/teowebrtc_log"
)

var log = teowebrtc_log.GetLog(teowebrtc_log.Package_teowebrtc_signal)

// New create new teonet webrtc signal server
func New(addr, addrSSL string, certAndKey ...string) (sig *Signal) {
	sig = new(Signal)
	sig.peers = make(map[string]*websocket.Conn)
	sig.serveWS(addr, addrSSL, certAndKey...)
	return
}

type Signal struct {
	peers map[string]*websocket.Conn
	sync.RWMutex
}

type loginSignal struct {
	Signal string `json:"signal"`
	Login  string `json:"login"`
}

type signalSignal struct {
	Signal string      `json:"signal"`
	Peer   string      `json:"peer"`
	Data   interface{} `json:"data"`
}

func (sig *Signal) addPeer(address string, conn *websocket.Conn) {
	sig.Lock()
	defer sig.Unlock()
	sig.peers[address] = conn
}

func (sig *Signal) delPeer(address string) {
	sig.Lock()
	defer sig.Unlock()
	delete(sig.peers, address)
}

// getPeer get peer connection by address
func (sig *Signal) getPeer(address string) (conn *websocket.Conn, ok bool) {
	sig.RLock()
	defer sig.RUnlock()
	conn, ok = sig.peers[address]
	return
}

// address get peer address by connection
func (sig *Signal) address(conn *websocket.Conn) string {
	sig.RLock()
	defer sig.RUnlock()
	for addr, c := range sig.peers {
		if c == conn {
			return addr
		}
	}
	return ""
}

// serveWS start websocek and websocket TLS server and listen for connections,
// the sertAndKey shoud contain certificate file path and private key file path
// to start HTTPS websocket server
func (sig *Signal) serveWS(addr, addrTLS string, certAndKey ...string) {
	http.HandleFunc("/signal", sig.signal)
	http.HandleFunc("/hello", sig.hello)
	http.HandleFunc("/", sig.home)
	go func() {
		log.Println("strat ws server at:", addr)
		log.Fatal(http.ListenAndServe(addr, nil))
	}()
	select {}
	log.Println("strat wws server at:", addrTLS)
	var certFile = "server.crt"
	var keyFile = "server.key"
	if len(certAndKey) > 0 {
		certFile = certAndKey[0]
	}
	if len(certAndKey) > 1 {
		keyFile = certAndKey[1]
	}
	log.Fatal(http.ListenAndServeTLS(addrTLS, certFile, keyFile, nil))
}

var upgrader = websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} // use default options

func (sig *Signal) signal(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade error:", err)
		return
	}
	defer c.Close()

	// Signals
	const (
		login     = "login"
		offer     = "offer"
		answer    = "answer"
		candidate = "candidate"
	)

	// This connection peer address
	var address string

	// writeErrMessage write json message with error
	writeErrMessage := func(messageType int, err error) {
		msg := []byte(fmt.Sprintf("{\"err\":\"%s\"}", err))
		log.Print("error: ", err)
		c.WriteMessage(messageType, msg)
	}

	// Check connected
	connected := func(messageType int) bool {
		if address != "" {
			return true
		}
		err := errors.New("should login first")
		writeErrMessage(messageType, err)
		return false
	}

	// Befor close
	defer func() { log.Printf("connection to %s closed", address) }()

	// Process messages
	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Println("read message error:", err)
			return
		}

		// Parse json
		var j map[string]interface{}
		json.Unmarshal(message, &j)

		// Parse signal
		signal, ok := j["signal"]
		if !ok {
			err := errors.New("wrong command, signal must be present")
			writeErrMessage(mt, err)
			return
		}

		switch signal {

		case login:
			// Unmarshal signal
			var s loginSignal
			err = json.Unmarshal(message, &s)
			if err != nil {
				writeErrMessage(mt, err)
				return
			}

			// Save login
			address = s.Login

			log.Printf("got %s from %s", signal, address)

			// Add peer to users map and remove it when connection closed
			sig.addPeer(s.Login, c)
			defer sig.delPeer(s.Login)

		case offer, answer, candidate:
			// Check login was done
			if !connected(mt) {
				return
			}

			log.Printf("got %s from %s", signal, address)

			// Unmarshal signal
			var s signalSignal
			err = json.Unmarshal(message, &s)
			if err != nil {
				writeErrMessage(mt, err)
				return
			}

			// Resend signal to peer
			log.Printf("resend %s to %s", s.Signal, s.Peer)
			conn, ok := sig.getPeer(s.Peer)
			if !ok {
				err := errors.New("peer " + s.Peer + " does not connected")
				writeErrMessage(mt, err)
				continue
			}
			s.Peer = address
			message, err = json.Marshal(s)
			if err != nil {
				log.Println("can't marshal")
				return
			}
			conn.WriteMessage(mt, message)
			continue
		}

		// Send answer
		err = c.WriteMessage(mt, message)
		if err != nil {
			log.Println("write message error:", err)
			return
		}
	}
}

func (sig *Signal) hello(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte("This is teonet webrtc signal server.\n"))
}

func (sig *Signal) home(w http.ResponseWriter, r *http.Request) {
	homeTemplate.Execute(w, "wss://"+r.Host+"/signal")
}

var homeTemplate = template.Must(template.New("").Parse(`
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<script>  
window.addEventListener("load", function(evt) {

    var output = document.getElementById("output");
    var input = document.getElementById("input");
    var ws;

    var print = function(message) {
        var d = document.createElement("div");
        d.textContent = message;
        output.appendChild(d);
        output.scroll(0, output.scrollHeight);
    };

    document.getElementById("open").onclick = function(evt) {
        if (ws) {
            return false;
        }
        ws = new WebSocket("{{.}}");
        ws.onopen = function(evt) {
            print("OPEN");
        }
        ws.onclose = function(evt) {
            print("CLOSE");
            ws = null;
        }
        ws.onmessage = function(evt) {
            print("RESPONSE: " + evt.data);
        }
        ws.onerror = function(evt) {
            print("ERROR: " + evt.data);
        }
        return false;
    };

    document.getElementById("send").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        print("SEND: " + input.value);
        ws.send(input.value);
        return false;
    };

    document.getElementById("close").onclick = function(evt) {
        if (!ws) {
            return false;
        }
        ws.close();
        return false;
    };

});
</script>
</head>
<body>
<table>
<tr><td valign="top" width="50%">
<h1>Signal server test</h1>
<p>Click "Open" to create a connection to the server, 
"Send" to send a message to the server and "Close" to close the connection. 
You can change the message and send multiple times.
<p>
<form>
<button id="open">Open</button>
<button id="close">Close</button>
<p><input id="input" type="text" value="Hello world!">
<button id="send">Send</button>
</form>
</td><td valign="top" width="50%">
<div id="output" style="max-height: 70vh;overflow-y: scroll;"></div>
</td></tr></table>
</body>
</html>
`))