Sharding router for Discord slash command webhooks

sharding-router.go

package main

import (
	"crypto/ed25519"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"os"
	"strconv"

	"github.com/valyala/fasthttp"
)

const (
	base_port = 3000
	shards    = uint64(4)
)

var (
	addr               = flag.String("addr", ":"+fmt.Sprint(base_port), "TCP address to listen to")
	public_key_data, _ = hex.DecodeString(os.Getenv("discordPublicKey"))
	public_key         = ed25519.PublicKey(public_key_data)
)

var proxyClients [shards]fasthttp.HostClient

func main() {
	flag.Parse()

	if len(public_key) != ed25519.PublicKeySize {
		log.Fatal("invalid discordPublicKey provided")
	}

	for i := range proxyClients {
		proxyClients[i] = fasthttp.HostClient{
			Addr: fmt.Sprintf("%s:%d", "127.0.0.1", base_port+1+i),
		}
	}

	if err := fasthttp.ListenAndServe(*addr, requestHandler); err != nil {
		log.Fatalf("Error in ListenAndServe: %s", err)
	}
}

type DiscordBody struct {
	GuildId string `json:"guild_id"`
}

func requestError(ctx *fasthttp.RequestCtx, code int, err string) {
	ctx.SetContentType("application/json; charset=utf8")
	fmt.Fprintf(ctx, "{\"error\": \"%s\"}", err)
	ctx.SetStatusCode(code)
}

func requestHandler(ctx *fasthttp.RequestCtx) {
	// Discord is only expected to send POST requests, so don't bother with anything else
	if !ctx.IsPost() {
		requestError(ctx, 400, "only POST is supported")
		return
	}

	// Verify signature
	body := ctx.PostBody()
	signature_data, err := hex.DecodeString(string(ctx.Request.Header.Peek("X-Signature-ed25519")))
	if err != nil {
		requestError(ctx, 500, "failed to decode signature")
		return
	}
	signed_data := append(ctx.Request.Header.Peek("X-Signature-Timestamp"), body...)
	if !ed25519.Verify(public_key, signed_data, signature_data) {
		requestError(ctx, 401, "invalid signature")
		return
	}

	// Send all trafic to the first shard by default
	shard := 0

	// Try decoding the body and direct trafic based on guild_id % shards
	var parsed DiscordBody
	if json.Unmarshal(body, &parsed) == nil {
		guild_id, err := strconv.ParseUint(parsed.GuildId, 10, 64)
		if err != nil {
			requestError(ctx, 500, "failed to parse guild_id")
			return
		}

		shard = int(guild_id % shards)
	}

	// Use a proxy client to forward the request
	err = proxyClients[shard].Do(&ctx.Request, &ctx.Response)
	if err != nil {
		requestError(ctx, 500, "failed to forward request")
		return
	}
}