keys/server/main.go

package main

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

	"github.com/gorilla/websocket"
	"github.com/joho/godotenv"
	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
	"github.com/labstack/gommon/log"
)

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 = ""
var upgrader = websocket.Upgrader{}

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

func main() {
	_ = godotenv.Load()
	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("/keys", keys)

	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 keys(c echo.Context) error {
	ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
	if err != nil {
		return err

	}
	defer ws.Close()

	client := Client{
		ID:   id,
		conn: ws,
		Exit: false,
	}
	clients = append(clients, &client)
	id += 1
	fmt.Printf("%s client %-3d connected\n", time.Now().Format(TimeFormat), client.ID)

	err = ws.WriteMessage(websocket.TextMessage, []byte("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)\n", time.Now().Format(TimeFormat), client.ID)
		client.Exit = true
		return c.NoContent(http.StatusInternalServerError)
	}
	defer f.Close()

	close := func() {
		_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		_ = ws.Close()
		client.Exit = true
	}

	for {
		if client.Exit {
			close()
			fmt.Printf("%s client %-3d disconnected\n", time.Now().Format(TimeFormat), client.ID)
			break
		}

		_, msg, err := ws.ReadMessage()
		if err != nil {
			log.Error(err)
			close()
			fmt.Printf("%s client %-3d crashed (websocket error)\n", time.Now().Format(TimeFormat), client.ID)
			break
		}

		_, err = f.Write(msg)
		if err != nil {
			log.Error(err)
			close()
			fmt.Printf("%s client %-3d crashed (file write error)\n", time.Now().Format(TimeFormat), client.ID)
			break
		}

	}

	return c.NoContent(http.StatusOK)
}