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 (
	"bytes"
	"context"
	"fmt"
	"io"
	"os"
	"sync"
	"time"

	"github.com/justincpresley/go-cobs"
	"github.com/shamaton/msgpackgen/msgpack"
	"github.com/sirupsen/logrus"

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

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

var cobsConfig = cobs.Config{SpecialByte: 0x00, Delimiter: true, EndingSave: true}

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()

	select {
	case 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.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 = 10 * time.Second

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

	time.Sleep(healthInterval)

	g.commands <- &protocol.Command{
		Announce: messages.BuildCommandAnnounce(messages.CmdFetchCRL),
		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:
			g.commands <- &protocol.Command{
				Announce: messages.BuildCommandAnnounce(messages.CmdHealth),
				Command:  &messages.HealthCommand{},
			}
		}

		timer.Reset(healthInterval)
	}
}

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

func (c *clientSimulator) readFrames() error {
	const readInterval = 50 * time.Millisecond

	var frame []byte
	buffer := &bytes.Buffer{}
	delimiter := []byte{cobsConfig.SpecialByte}

	for {
		readBytes, err := c.readFromStdin()
		if err != nil {
			c.logger.WithError(err).Error("stdin read error")

			close(c.framesIn)

			return err
		}

		if len(readBytes) == 0 {
			time.Sleep(readInterval)

			continue
		}

		c.logger.Tracef("read %d bytes", len(readBytes))

		buffer.Write(readBytes)

		c.logger.Tracef("read buffer is now %d bytes long", buffer.Len())

		rest := buffer.Bytes()

		if !bytes.Contains(rest, delimiter) {
			continue
		}

		for bytes.Contains(rest, delimiter) {
			parts := bytes.SplitAfterN(rest, delimiter, 2)
			frame, rest = parts[0], parts[1]

			c.logger.Tracef("frame of length %d", len(frame))

			if len(frame) == 0 {
				continue
			}

			err = cobs.Verify(frame, cobsConfig)
			if err != nil {
				return fmt.Errorf("frame verification failed: %w", err)
			}

			decoded := cobs.Decode(frame, cobsConfig)

			c.logger.Tracef("frame decoded to length %d", len(decoded))

			c.framesIn <- decoded
		}

		buffer.Truncate(0)
		buffer.Write(rest)

		c.logger.Tracef("read buffer is now %d bytes long", buffer.Len())
	}
}

func (c *clientSimulator) writeFrame(frame []byte) error {
	encoded := cobs.Encode(frame, cobsConfig)

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

	if _, err := io.Copy(os.Stdout, bytes.NewBuffer(encoded)); err != nil {
		return fmt.Errorf("could not write data: %w", err)
	}

	return nil
}

func (c *clientSimulator) readFromStdin() ([]byte, error) {
	const bufferSize = 1024

	buf := make([]byte, bufferSize)

	count, err := os.Stdin.Read(buf)
	if err != nil {
		return nil, fmt.Errorf("reading input failed: %w", err)
	}

	return buf[:count], nil
}

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")

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

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

	return nil
}

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

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

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

	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")

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

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

	return nil
}

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

	select {
	case 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")

	select {
	case 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.readFrames()

		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)

	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.TraceLevel)

	messages.RegisterGeneratedResolver()

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

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