cacert-gosigner/cmd/clientsim/main.go

/*
Copyright 2022 CAcert Inc.
SPDX-License-Identifier: Apache-2.0

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// client simulator

package main

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/shamaton/msgpackgen/msgpack"
	"github.com/sirupsen/logrus"

	"git.cacert.org/cacert-gosigner/pkg/protocol"

	"git.cacert.org/cacert-gosigner/pkg/messages"
)

type protocolState int8

const (
	cmdAnnounce protocolState = iota
	cmdData
	respAnnounce
	respData
)

var validTransitions = map[protocolState]protocolState{
	cmdAnnounce:  cmdData,
	cmdData:      respAnnounce,
	respAnnounce: respData,
	respData:     cmdAnnounce,
}

var protocolStateNames = map[protocolState]string{
	cmdAnnounce:  "CMD ANNOUNCE",
	cmdData:      "CMD DATA",
	respAnnounce: "RESP ANNOUNCE",
	respData:     "RESP DATA",
}

func (p protocolState) String() string {
	if name, ok := protocolStateNames[p]; ok {
		return name
	}

	return fmt.Sprintf("unknown %d", p)
}

type TestCommandGenerator struct {
	logger          *logrus.Logger
	currentCommand  *protocol.Command
	currentResponse *protocol.Response
	commands        chan *protocol.Command
	lock            sync.Mutex
}

func (g *TestCommandGenerator) CmdAnnouncement() ([]byte, error) {
	g.lock.Lock()
	defer g.lock.Unlock()

	g.currentCommand = <-g.commands

	announceData, err := msgpack.Marshal(g.currentCommand.Announce)
	if err != nil {
		return nil, fmt.Errorf("could not marshal command annoucement: %w", err)
	}

	g.logger.WithField("announcement", &g.currentCommand.Announce).Info("write command announcement")

	return announceData, nil
}

func (g *TestCommandGenerator) CmdData() ([]byte, error) {
	g.lock.Lock()
	defer g.lock.Unlock()

	cmdData, err := msgpack.Marshal(g.currentCommand.Command)
	if err != nil {
		return nil, fmt.Errorf("could not marshal command data: %w", err)
	}

	g.logger.WithField("command", &g.currentCommand.Command).Info("write command data")

	return cmdData, nil
}

func (g *TestCommandGenerator) HandleResponseAnnounce(frame []byte) error {
	g.lock.Lock()
	defer g.lock.Unlock()

	var ann messages.ResponseAnnounce

	if err := msgpack.Unmarshal(frame, &ann); err != nil {
		return fmt.Errorf("could not unmarshal response announcement")
	}

	g.logger.WithField("announcement", &ann).Info("received response announcement")

	g.currentResponse = &protocol.Response{Announce: &ann}

	return nil
}

func (g *TestCommandGenerator) HandleResponse(frame []byte) error {
	g.lock.Lock()
	defer g.lock.Unlock()

	switch g.currentResponse.Announce.Code {
	case messages.RespHealth:
		var response messages.HealthResponse

		if err := msgpack.Unmarshal(frame, &response); err != nil {
			return fmt.Errorf("unmarshal failed: %w", err)
		}

		g.currentResponse.Response = &response
	case messages.RespFetchCRL:
		var response messages.FetchCRLResponse

		if err := msgpack.Unmarshal(frame, &response); err != nil {
			return fmt.Errorf("unmarshal failed: %w", err)
		}

		g.currentResponse.Response = &response
	}

	g.logger.WithField(
		"command",
		g.currentCommand,
	).WithField(
		"response",
		g.currentResponse,
	).Info("handled health response")

	return nil
}

func (g *TestCommandGenerator) GenerateCommands(ctx context.Context) error {
	const (
		healthInterval = 5 * time.Second
		startPause     = 3 * time.Second
	)

	var (
		announce *messages.CommandAnnounce
		err      error
	)

	g.logger.Info("start generating commands")

	time.Sleep(startPause)

	announce, err = messages.BuildCommandAnnounce(messages.CmdFetchCRL)
	if err != nil {
		return fmt.Errorf("build command announce failed: %w", err)
	}

	g.commands <- &protocol.Command{
		Announce: announce,
		Command:  &messages.FetchCRLCommand{IssuerID: "sub-ecc_person_2022"},
	}

	timer := time.NewTimer(healthInterval)

	for {
		select {
		case <-ctx.Done():
			_ = timer.Stop()

			g.logger.Info("stopping health check loop")

			return nil
		case <-timer.C:
			announce, err = messages.BuildCommandAnnounce(messages.CmdHealth)
			if err != nil {
				return fmt.Errorf("build command announce failed: %w", err)
			}

			g.commands <- &protocol.Command{
				Announce: announce,
				Command:  &messages.HealthCommand{},
			}
		}

		timer.Reset(healthInterval)
	}
}

type clientSimulator struct {
	protocolState    protocolState
	logger           *logrus.Logger
	lock             sync.Mutex
	framesIn         chan []byte
	framesOut        chan []byte
	framer           protocol.Framer
	commandGenerator *TestCommandGenerator
}

func (c *clientSimulator) writeCmdAnnouncement() error {
	frame, err := c.commandGenerator.CmdAnnouncement()
	if err != nil {
		return fmt.Errorf("could not get command annoucement: %w", err)
	}

	c.logger.Trace("writing command announcement")

	c.framesOut <- frame

	if err := c.nextState(); err != nil {
		return err
	}

	return nil
}

func (c *clientSimulator) writeCommand() error {
	frame, err := c.commandGenerator.CmdData()
	if err != nil {
		return fmt.Errorf("could not get command data: %w", err)
	}

	c.logger.Trace("writing command data")

	c.framesOut <- frame

	if err := c.nextState(); err != nil {
		return err
	}

	return nil
}

func (c *clientSimulator) handleResponseAnnounce() error {
	c.logger.Trace("waiting for response announce")

	frame := <-c.framesIn

	if frame == nil {
		return nil
	}

	if err := c.commandGenerator.HandleResponseAnnounce(frame); err != nil {
		return fmt.Errorf("response announce handling failed: %w", err)
	}

	if err := c.nextState(); err != nil {
		return err
	}

	return nil
}

func (c *clientSimulator) handleResponseData() error {
	c.logger.Trace("waiting for response data")

	frame := <-c.framesIn

	if frame == nil {
		return nil
	}

	if err := c.commandGenerator.HandleResponse(frame); err != nil {
		return fmt.Errorf("response handler failed: %w", err)
	}

	if err := c.nextState(); err != nil {
		return err
	}

	return nil
}

func (c *clientSimulator) Run(ctx context.Context) error {
	c.protocolState = cmdAnnounce
	errors := make(chan error)

	go func() {
		err := c.framer.ReadFrames(os.Stdin, c.framesIn)

		errors <- err
	}()

	go func() {
		err := c.framer.WriteFrames(os.Stdout, c.framesOut)

		errors <- err
	}()

	go func() {
		err := c.commandGenerator.GenerateCommands(ctx)

		errors <- err
	}()

	for {
		select {
		case <-ctx.Done():
			return nil
		case err := <-errors:
			if err != nil {
				return fmt.Errorf("error from handler loop: %w", err)
			}

			return nil
		default:
			if err := c.handleProtocolState(); err != nil {
				return err
			}
		}
	}
}

func (c *clientSimulator) handleProtocolState() error {
	c.logger.Tracef("handling protocol state %s", c.protocolState)

	c.lock.Lock()
	defer c.lock.Unlock()

	switch c.protocolState {
	case cmdAnnounce:
		if err := c.writeCmdAnnouncement(); err != nil {
			return err
		}
	case cmdData:
		if err := c.writeCommand(); err != nil {
			return err
		}
	case respAnnounce:
		if err := c.handleResponseAnnounce(); err != nil {
			return err
		}
	case respData:
		if err := c.handleResponseData(); err != nil {
			return err
		}
	default:
		return fmt.Errorf("unknown protocol state %s", c.protocolState)
	}

	return nil
}

func (c *clientSimulator) nextState() error {
	next, ok := validTransitions[c.protocolState]
	if !ok {
		return fmt.Errorf("illegal protocol state %s", c.protocolState)
	}

	c.protocolState = next

	return nil
}

func main() {
	logger := logrus.New()
	logger.SetOutput(os.Stderr)
	logger.SetLevel(logrus.DebugLevel)

	messages.RegisterGeneratedResolver()

	sim := &clientSimulator{
		commandGenerator: &TestCommandGenerator{
			logger:   logger,
			commands: make(chan *protocol.Command),
		},
		logger:    logger,
		framesIn:  make(chan []byte),
		framesOut: make(chan []byte),
		framer:    protocol.NewCOBSFramer(logger),
	}

	err := sim.Run(context.Background())
	if err != nil {
		logger.WithError(err).Error("simulator returned an error")
	}
}