go_sshkeys/internal/utils.go

package internal

import (
	"bytes"
	"encoding/binary"
	`io`

	`r00t2.io/sshkeys/errs`
)

/*
	ReadSizeBytes returns a uint32 allocator for a given set of data for preparation of packing into bytes.

	See SerializeData for valid types for data.

	If pack is true, the original data will be appended to the returned allocated *bytes.Reader.

	NOTE: If data is a *bytes.Reader, the bytes WILL be consumed within that reader.
*/
func ReadSizeBytes(data interface{}, pack bool) (allocated *bytes.Reader, err error) {

	var u uint32
	var b []byte
	var sizer = make([]byte, 4)

	if b, err = SerializeData(data); err != nil {
		return
	}

	u = uint32(len(b))

	if !pack {
		binary.BigEndian.PutUint32(sizer, u)
		allocated = bytes.NewReader(sizer)
	} else {
		if b, err = PackBytes(b); err != nil {
			return
		}
		allocated = bytes.NewReader(b)
	}

	return
}

/*
	PackBytes prefixes data with an allocator.

	See SerializeData for valid types for data.

	NOTE: If data is a *bytes.Reader, the bytes WILL be consumed within that reader.
*/
func PackBytes(data interface{}) (b []byte, err error) {

	var dataBytes []byte
	var size uint32

	if dataBytes, err = SerializeData(data); err != nil {
		return
	}

	size = uint32(len(dataBytes))

	b = make([]byte, size)
	binary.BigEndian.PutUint32(b, size)
	b = append(b, dataBytes...)

	return
}

/*
	UnpackBytes performs the opposite of PackBytes, reading an allocator and then reading ONLY the remaining bytes.

	NOTE: If data is a bytes.Buffer pointer, it will consume ONLY the leading prefix and the length of bytes the prefix indicates and no more.

	NOTE: If data is a *bytes.Reader, it will consume ONLY the leading prefix and the length of bytes the prefix indicates and no more.
*/
func UnpackBytes(data interface{}) (unpacked []byte, err error) {

	var sizeBytes []byte = make([]byte, 4)
	var u uint32

	// We can't use SerializeData because it'll consume the entire data set. We don't want that.
	switch t := data.(type) {
	case []byte:
		sizeBytes = t[:4]
		u = binary.BigEndian.Uint32(sizeBytes)
		unpacked = t[4:u]
	case *bytes.Buffer:
		sizeBytes = t.Next(4)
		u = binary.BigEndian.Uint32(sizeBytes)
		unpacked = t.Next(int(u))
	case bytes.Buffer:
		t2 := &t
		sizeBytes = t2.Next(4)
		u = binary.BigEndian.Uint32(sizeBytes)
		unpacked = t2.Next(int(u))
	case *bytes.Reader:
		if _, err = t.Read(sizeBytes); err != nil {
			return
		}
		u = binary.BigEndian.Uint32(sizeBytes)
		unpacked = make([]byte, u)
		if _, err = io.ReadFull(t, unpacked); err != nil {
			return
		}
	case bytes.Reader:
		t2 := &t
		if _, err = t2.Read(sizeBytes); err != nil {
			return
		}
		u = binary.BigEndian.Uint32(sizeBytes)
		unpacked = make([]byte, u)
		if _, err = io.ReadFull(t2, unpacked); err != nil {
			return
		}
	default:
		err = errs.ErrBadData
		return
	}

	return
}

/*
	SerializeData transforms one of the following types:

		- string
		- []byte
		- byte
		- bytes.Buffer
		- *bytes.Buffer
		- bytes.Reader
		- *bytes.Reader

	into a []byte.

	NOTE: If data is a *bytes.Buffer, no bytes will be consumed -- the bytes are taken in entirety without consuming them (Buffer.Bytes()).

	NOTE: If data is a *bytes.Reader, ALL bytes WILL be consumed.
*/
func SerializeData(data interface{}) (b []byte, err error) {

	switch t := data.(type) {
	case string:
		b = []byte(t)
	case byte:
		b = []byte{t}
	case []byte:
		b = t
	case *bytes.Buffer:
		b = t.Bytes()
	case bytes.Buffer:
		t2 := &t
		b = t2.Bytes()
	case *bytes.Reader:
		b = make([]byte, t.Len())
		if b, err = io.ReadAll(t); err != nil {
			return
		}
	case bytes.Reader:
		t2 := &t
		b = make([]byte, t2.Len())
		if b, err = io.ReadAll(t2); err != nil {
			return
		}
	default:
		err = errs.ErrBadData
		return
	}

	return
}