ngkrok/internal/ngkrok/server.go

119 lines
2.3 KiB
Go

package ngkrok
import (
"context"
"crypto/rand"
"io"
"log"
"net"
"net/http"
"strings"
vhost "github.com/inconshreveable/go-vhost"
"github.com/progrium/qmux/golang/session"
)
const (
SERVER_VERSION = "v0.5"
HEADER_CLIENT_TOKEN = "X-Client-App-Token"
HEADER_PUBLIC_HOST = "X-Public-Host"
)
type Server struct {
addr string
host string
port string
VMux *vhost.HTTPMuxer
}
func NewServer(addr, host, port string, vmux *vhost.HTTPMuxer) *Server {
return &Server{addr, host, port, vmux}
}
func (s *Server) Serve() {
ml, err := s.VMux.Listen(net.JoinHostPort(s.host, s.port))
if err != nil {
log.Fatal(err)
}
srv := &http.Server{}
// http.HandleFunc("/", s.mainHandler)
http.HandleFunc("/register", s.registerHandler)
srv.Serve(ml)
}
// func (s *Server) mainHandler(w http.ResponseWriter, r *http.Request) {
// log.Println("mainHandler")
// return
// }
func (s *Server) registerHandler(w http.ResponseWriter, r *http.Request) {
subdomain := newSubdomain(r) + s.host
publicHost := strings.TrimSuffix(net.JoinHostPort(subdomain, s.port), ":80")
pl, err := s.VMux.Listen(publicHost)
if err != nil {
log.Fatal(err)
}
defer pl.Close()
w.Header().Add(HEADER_PUBLIC_HOST, subdomain)
w.Header().Add("Connection", "close")
w.WriteHeader(http.StatusOK)
conn, _, _ := w.(http.Hijacker).Hijack()
sess := session.New(conn)
defer sess.Close()
log.Printf("%s: start session", subdomain)
go func() {
for {
conn, err := pl.Accept()
if err != nil {
log.Println("Accept:", err)
return
}
defer conn.Close()
ch, err := sess.Open(context.Background())
if err != nil {
log.Println("sess.Open:", err)
return
}
defer ch.Close()
go join(ch, conn)
}
}()
sess.Wait()
log.Printf("%s: end session", subdomain)
}
func join(a io.ReadWriteCloser, b io.ReadWriteCloser) {
go io.Copy(b, a)
io.Copy(a, b)
a.Close()
b.Close()
}
func newSubdomain(r *http.Request) string {
clientToken := r.Header.Get(HEADER_CLIENT_TOKEN)
if clientToken == "ghd7H8H*&6tg8*tg68" { // hardcoded token for dev tests....
return "test."
}
b := make([]byte, 10)
if _, err := rand.Read(b); err != nil {
panic(err)
}
letters := []rune("abcdefghijklmnopqrstuvwxyz1234567890")
rune := make([]rune, 10)
for i := range rune {
rune[i] = letters[int(b[i])*len(letters)/256]
}
return string(rune) + "."
}