iconia/cmd/iconiad/main.go

package main

import (
	"context"
	"crypto/tls"
	"flag"
	"net/http"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/facebookarchive/flagenv"
	"github.com/hashicorp/yamux"
	"go.chromium.org/luci/common/flag/stringmapflag"
	"golang.org/x/crypto/acme/autocert"
	"within.website/confyg/flagconfyg"
	"within.website/ln"
	"within.website/x/localca"
)

var (
	httpPort   = flag.String("http-port", "3043", "HTTP port")
	httpsPort  = flag.String("https-port", "3044", "HTTPS port")
	yamuxPort  = flag.String("yamux-port", "3045", "yamux port")
	statusPort = flag.String("status-port", "3046", "status server port")

	// TLS certificate configuration
	domainSuffix = flag.String("domain-suffix", ".local.cetacean.club", "allowed domain suffix for certificate generation")
	certFile     = flag.String("cert-file", "./var/minica.pem", "TLS certificate authority public certificate")
	keyFile      = flag.String("key-file", "./var/minica-key.pem", "TLS certificate authority private key")
	certFolder   = flag.String("cert-folder", "./var/certs", "TLS certificate storage folder")

	// hosts -> tokens
	hostsToTokens = new(stringmapflag.Value)
)

func init() {
	flag.Var(hostsToTokens, "host-token", "accepted pairs of hostname -> token mappings")
}

func main() {
	flagenv.Parse()
	flagconfyg.CmdParse("./iconia.confyg")
	flag.Parse()

	cfg := Config{
		HTTPPort:     *httpPort,
		HTTPSPort:    *httpsPort,
		YamuxPort:    *yamuxPort,
		StatusPort:   *statusPort,
		DomainSuffix: *domainSuffix,
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ctx = ln.WithF(ctx, cfg.F())

	certManager, err := localca.New(*keyFile, *certFile, *domainSuffix, autocert.DirCache(*certFolder))
	if err != nil {
		ln.FatalErr(ctx, err)
	}

	httpsTc := &tls.Config{
		GetCertificate: certManager.GetCertificate,
	}

	httpsListener, err := tls.Listen("tcp", ":"+*httpsPort, httpsTc)
	if err != nil {
		ln.FatalErr(ctx, err)
	}

	s := &Server{
		Config: cfg,

		clients:     map[string][]*yamux.Session{},
		clientsLock: &sync.RWMutex{},

		tokenInfo:  map[string]string(*hostsToTokens),
		tokensLock: &sync.Mutex{},

		certManager: certManager,

		tlsListener: httpsListener,

		plainServer: &http.Server{
			Addr:    ":" + *httpPort,
			Handler: http.HandlerFunc(insecureRedirect),
		},

		statusServer: &http.Server{
			Addr:    ":" + *statusPort,
			Handler: http.DefaultServeMux,
		},
	}
	_ = s

	yamuxTc := &tls.Config{
		GetCertificate:     certManager.GetCertificate,
		GetConfigForClient: s.handleYamuxClientHello,
	}

	yamuxListener, err := tls.Listen("tcp", ":"+*yamuxPort, yamuxTc)
	if err != nil {
		ln.FatalErr(ctx, err)
	}
	s.yamuxListener = yamuxListener

	ln.Log(ctx, ln.Info("now listening for traffic"))

	go func() { ln.Error(ctx, s.plainServer.ListenAndServe()) }()
	go func() { ln.Error(ctx, s.statusServer.ListenAndServe()) }()
	go func() { ln.Error(ctx, s.tlsForward(httpsListener)) }()
	go func() { ln.Error(ctx, s.yamuxHandler(yamuxListener)) }()

	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)

	<-sigs
	ln.Log(ctx, ln.Info("got SIGINT/SIGTERM, dying"))
	defer cancel()

	s.yamuxListener.Close()
	s.tlsListener.Close()
	s.plainServer.Shutdown(ctx)
	s.statusServer.Shutdown(ctx)
	s.goAwayClients()

	a := time.After(4 * time.Minute)
	select {
	case <-a:
		s.killClients()
	case <-sigs:
		s.killClients()
	}
}

// insecureRedirect redirects a client to https if they connect over plain HTTP.
func insecureRedirect(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodPatch, http.MethodPut, http.MethodPost:
		http.Error(w, "use https", http.StatusNotAcceptable)
		ln.Log(r.Context(), ln.Action("cannot redirect (wrong method)"), ln.F{"remote": r.RemoteAddr, "host": r.Host, "path": r.URL.Path})
		return
	}

	r.URL.Host = r.Host
	r.URL.Scheme = "https"

	ln.Log(r.Context(), ln.Action("redirecting insecure HTTP to HTTPS"), ln.F{"remote": r.RemoteAddr, "host": r.Host, "path": r.URL.Path})

	http.Redirect(w, r, r.URL.String(), http.StatusPermanentRedirect)
}