lnd/watchtower/wtwire/fuzz_test.go

package wtwire

import (
	"bytes"
	"encoding/binary"
	"testing"

	"github.com/stretchr/testify/require"
)

// prefixWithMsgType takes []byte and adds a wire protocol prefix
// to make the []byte into an actual message to be used in fuzzing.
func prefixWithMsgType(data []byte, prefix MessageType) []byte {
	var prefixBytes [2]byte
	binary.BigEndian.PutUint16(prefixBytes[:], uint16(prefix))
	data = append(prefixBytes[:], data...)

	return data
}

// wireMsgHarness performs the actual fuzz testing of the appropriate wire
// message. This function will check that the passed-in message passes wire
// length checks, is a valid message once deserialized, and passes a sequence of
// serialization and deserialization checks. emptyMsg must be an empty Message
// of the type to be fuzzed, as it is used to determine the appropriate prefix
// bytes and max payload length for decoding.
func wireMsgHarness(t *testing.T, data []byte, emptyMsg Message) {
	t.Helper()

	// Make sure byte array length is less than max payload size for the
	// wire message.
	payloadLen := uint32(len(data))
	if payloadLen > emptyMsg.MaxPayloadLength(0) {
		// Ignore this input - max payload constraint violated.
		return
	}

	data = prefixWithMsgType(data, emptyMsg.MsgType())

	// Create a reader with the byte array.
	r := bytes.NewReader(data)

	msg, err := ReadMessage(r, 0)
	if err != nil {
		return
	}

	// We will serialize the message into a new bytes buffer.
	var b bytes.Buffer
	_, err = WriteMessage(&b, msg, 0)
	require.NoError(t, err)

	// Deserialize the message from the serialized bytes buffer, and then
	// assert that the original message is equal to the newly deserialized
	// message.
	newMsg, err := ReadMessage(&b, 0)
	require.NoError(t, err)
	require.Equal(t, msg, newMsg)
}

func FuzzCreateSessionReply(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &CreateSessionReply{})
	})
}

func FuzzCreateSession(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &CreateSession{})
	})
}

func FuzzDeleteSessionReply(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &DeleteSessionReply{})
	})
}

func FuzzDeleteSession(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &DeleteSession{})
	})
}

func FuzzError(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &Error{})
	})
}

func FuzzInit(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &Init{})
	})
}

func FuzzStateUpdateReply(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &StateUpdateReply{})
	})
}

func FuzzStateUpdate(f *testing.F) {
	f.Fuzz(func(t *testing.T, data []byte) {
		wireMsgHarness(t, data, &StateUpdate{})
	})
}