// This is an example message reader/writer with length prefixes and checksums
// (in the example's case, using hmac-sha1) to try to confirm that the entire
// correct message was read.

package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/binary"
	"errors"
	"fmt"
	"hash"
	"io"
)

func main() {
	var b bytes.Buffer

	var hasher HashFunc
	// CRC64:
	// isoTable := crc64.MakeTable(crc64.ISO)
	// hasher = HashFunc(func() hash.Hash { return crc64.New(isoTable) })
	// HMAC-SHA1:
	hasher = HMACFunc("key 1", sha1.New)
	// SHA1:
	// hasher = HashFunc(sha1.New)
	// No hash:
	// hasher = HashFunc(NullHash)

	err := NewMessenger(100, binary.BigEndian, hasher).WriteMsg(&b, []byte("Hello, World"))
	if err != nil {
		panic(err)
	}

	// Print written message
	fmt.Printf("BUF = %x\n", &b)

	// Create hash mismatch:
	// hasher = hmacHashFunc("key 2", sha1.New)

	msg, err := NewMessenger(100, binary.BigEndian, hasher).ReadMsg(&b)
	if err != nil {
		panic(err)
	}

	// Print read message
	fmt.Printf("MSG = %q\n", msg)
}

// HashFunc is a function that returns a new hash.Hash for use in a Messenger.
// Returned hashes are not reused.
type HashFunc func() hash.Hash

// Messenger reads and writes messages with length and checksum prefixes.
type Messenger struct {
	hashfn  HashFunc
	order   binary.ByteOrder
	maxSize int64 // TODO: Consider using varints to avoid wasting a whole 8 bytes
}

// NewMessenger allocates a new Messenger for reading and writing length-and-checksum prefixed
// messages. maxMsgSize specifies the maximum size of a message in bytes, and must be greater than
// eight bytes. If any arguments are invalid, NewMessenger panics.
//
// Returning an error would be more appropriate in real code, but this is a prototype.
func NewMessenger(maxMsgSize int64, order binary.ByteOrder, hashfn HashFunc) *Messenger {
	if order == nil {
		panic("nil byte order")
	}
	if hashfn == nil {
		panic("nil hash constructor")
	}
	if maxMsgSize <= 8 {
		panic("max message size must be > 8 bytes")
	}
	return &Messenger{
		hashfn:  hashfn,
		order:   order,
		maxSize: maxMsgSize,
	}
}

type msgHeader struct {
	Size     int64
	Checksum []byte // Includes Size's bytes and the payload in it
}

// ReadMsg reads length, checksum, and payload from r.
// If the checksum of the payload does not match when filtered through Messenger's hash function, it
// returns an error. This is to prevent acceptance of corrupt, spoofed, or otherwise invalid
// messages (depending on the checksum).
// In addition, too-small and too-large messages will also return errors.
// All other errors arise from reading from r.
func (m *Messenger) ReadMsg(r io.Reader) ([]byte, error) {
	h := m.hashfn()
	hashSize := h.Size()
	var header msgHeader

	hashReader := io.TeeReader(r, h)
	if err := binary.Read(hashReader, m.order, &header.Size); err != nil {
		return nil, err
	}

	headerSize := int64(hashSize)
	if header.Size > m.maxSize {
		return nil, fmt.Errorf("message of %d bytes exceeds max size of %d bytes", header.Size, m.maxSize)
	} else if header.Size <= headerSize {
		return nil, fmt.Errorf("message of %d bytes is too short", header.Size)
	}

	header.Checksum = make([]byte, hashSize)
	if _, err := io.ReadFull(r, header.Checksum); err != nil {
		return nil, err
	}

	p := make([]byte, int(header.Size-headerSize))
	_, err := io.ReadFull(hashReader, p)
	if err != nil {
		return nil, err
	}

	if sum := h.Sum(nil); !hmac.Equal(sum, header.Checksum) {
		return nil, fmt.Errorf("message checksums don't match: sent(%x) <> received(%x)", sum, header.Checksum)
	}

	return p, nil
}

// WriteMsg writes the payload, p, prefixed by length and checksum, to the writer w.
// If the length of p combined with the header's length would exceed the Messenger's maximum message
// size, then the message is not written and an error is returned. Empty messages also return an
// error. All other errors arise from writing to w.
func (m *Messenger) WriteMsg(w io.Writer, p []byte) error {
	if len(p) == 0 {
		return errors.New("message is empty")
	}

	h := m.hashfn()
	hashSize := h.Size()
	headerSize := int64(hashSize)

	if int64(len(p)) > m.maxSize-headerSize {
		return fmt.Errorf("message of %d bytes exceeds max size of %d - header(%d) bytes", len(p), m.maxSize, headerSize)
	}

	header := msgHeader{
		Size:     int64(len(p)) + headerSize,
		Checksum: make([]byte, hashSize),
	}

	if err := binary.Write(h, m.order, header.Size); err != nil {
		return err // This is probably unreachable.
	}
	h.Write(p)
	header.Checksum = h.Sum(header.Checksum[:0])

	if err := binary.Write(w, m.order, header.Size); err != nil {
		return err
	}

	if n, err := w.Write(header.Checksum); err != nil {
		return err
	} else if n != len(header.Checksum) {
		// NB: This doesn't handle the strange case of n > len() politely
		return io.ErrShortWrite
	}

	if n, err := w.Write(p); err != nil {
		return err
	} else if n != len(p) {
		return io.ErrShortWrite
	}

	return nil
}

// Hash functions for testing:

func HMACFunc(key string, hashfn HashFunc) HashFunc {
	bkey := []byte(key)
	return func() hash.Hash {
		return hmac.New(hashfn, bkey)
	}
}

func NullHash() hash.Hash {
	return nullHasher{}
}

type nullHasher struct{}

func (nullHasher) Write(p []byte) (n int, err error) {
	return len(p), nil
}

func (nullHasher) Size() int           { return 0 }
func (nullHasher) Reset()              {}
func (nullHasher) BlockSize() int      { return 1 }
func (nullHasher) Sum(h []byte) []byte { return h[:0] }