protohackers/cmd/primetime/main.go

package main

import (
	"encoding/json"
	"log"
	"math/big"
	"net"

	"protohackers/pkg/conn"
)

var (
	malformedResponse = []byte("malformed\n")
	trueResponse      = []byte(`{"method":"isPrime","prime":true}` + "\n")
	falseResponse     = []byte(`{"method":"isPrime","prime":false}` + "\n")
)

type data struct {
	Method string   `json:"method"`
	Number *float64 `json:"number"`
}

func main() {
	err := conn.StartSimpleTCP(primetime)
	if err != nil {
		log.Fatalln(err)
	}
}

func primetime(c net.Conn) {
	if err := conn.ReadAndRespond(c, process); err != nil {
		log.Fatalln(err)
	}
}

func process(buf []byte) []byte {
	var in data
	err := json.Unmarshal(buf, &in)
	if err != nil || in.Method != "isPrime" || in.Number == nil {
		return malformedResponse
	} else if big.NewInt(int64(*in.Number)).ProbablyPrime(0) {
		return trueResponse
	}
	return falseResponse
}