diff --git a/cmd/primetime/main.go b/cmd/primetime/main.go new file mode 100644 index 0000000..5921b0e --- /dev/null +++ b/cmd/primetime/main.go @@ -0,0 +1,52 @@ +package main + +import ( + "bufio" + "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.StartSimple(primetime) + if err != nil { + log.Fatalln(err) + } +} + +func primetime(c net.Conn) { + defer c.Close() + sc := bufio.NewScanner(c) + for sc.Scan() { + response := process(sc.Bytes()) + c.Write(response) + } + if err := sc.Err(); 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 +} diff --git a/cmd/smoketest/main.go b/cmd/smoketest/main.go index 10f924b..717bb7b 100644 --- a/cmd/smoketest/main.go +++ b/cmd/smoketest/main.go @@ -1,40 +1,23 @@ package main import ( - "fmt" "io" "log" "net" - "os" -) -const ( - connPort = "3030" - connType = "tcp4" + "protohackers/pkg/conn" ) func main() { - l, err := net.Listen(connType, ":"+connPort) + err := conn.StartSimple(smoketest) if err != nil { log.Fatalln(err) - os.Exit(1) - } - defer l.Close() - - fmt.Println(fmt.Sprintf("Listening on port %s", connPort)) - - for { - conn, err := l.Accept() - if err != nil { - log.Fatalln(err) - } - go handleRequest(conn) } } -func handleRequest(conn net.Conn) { - defer conn.Close() - if _, err := io.Copy(conn, conn); err != nil { +func smoketest(c net.Conn) { + defer c.Close() + if _, err := io.Copy(c, c); err != nil { log.Fatalln(err) } } diff --git a/pkg/conn/conn.go b/pkg/conn/conn.go new file mode 100644 index 0000000..9d266e9 --- /dev/null +++ b/pkg/conn/conn.go @@ -0,0 +1,32 @@ +package conn + +import ( + "fmt" + "net" +) + +const ( + Port = ":3030" + Type = "tcp4" +) + +func StartSimple(handler func(conn net.Conn)) error { + l, err := net.Listen(Type, Port) + if err != nil { + return err + } + defer l.Close() + + fmt.Println(fmt.Sprintf("Listening on port %s", Port)) + + for { + c, err := l.Accept() + if err != nil { + return err + } + go handler(c) + if err != nil { + return err + } + } +}