go_sshkeys/internal/utils.go

package internal

import (
	"bytes"
	"encoding/binary"
)

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

	data can be one of:

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

	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)

	switch t := data.(type) {
	case string:
		b = []byte(t)
	case byte:
		// b = []byte{0x0, 0x0, 0x0, 0x1}
		b = []byte{t}
	case []byte:
		b = t
	case bytes.Buffer:
		t2 := &t
		b = t2.Bytes()
	case *bytes.Buffer:
		b = t.Bytes()
	case bytes.Reader:
		t2 := &t
		b = make([]byte, t2.Len())
		if _, err = t2.Read(b); err != nil {
			return
		}
	case *bytes.Reader:
		b = make([]byte, t.Len())
		if _, err = t.Read(b); err != nil {
			return
		}
	}

	u = uint32(len(b))

	binary.BigEndian.PutUint32(sizer, u)
	if !pack {
		allocated = bytes.NewReader(sizer)
	} else {
		b = append(sizer, b...)
		allocated = bytes.NewReader(b)
	}

	return
}