feat: update

This commit is contained in:
sujit
2024-10-14 05:54:01 +05:45
parent d910a2656b
commit a5bc632b1e
10 changed files with 92 additions and 241 deletions

View File

@@ -1,32 +1,26 @@
package codec
import (
"bytes"
"crypto/aes"
"crypto/cipher"
"crypto/hmac"
"crypto/rand"
"crypto/sha256"
"encoding/binary"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net"
"sync"
"github.com/vmihailenco/msgpack/v5"
"github.com/oarkflow/mq/consts"
"github.com/oarkflow/mq/utils"
)
// Message represents the structure of our message.
type Message struct {
Headers map[string]string `json:"h"`
Queue string `json:"q"`
Command consts.CMD `json:"c"`
Payload json.RawMessage `json:"p"`
Headers map[string]string `msgpack:"h"`
Queue string `msgpack:"q"`
Command consts.CMD `msgpack:"c"`
Payload []byte `msgpack:"p"` // Using []byte instead of json.RawMessage
m sync.RWMutex
}
func NewMessage(cmd consts.CMD, payload json.RawMessage, queue string, headers map[string]string) *Message {
// NewMessage creates a new Message instance.
func NewMessage(cmd consts.CMD, payload []byte, queue string, headers map[string]string) *Message {
return &Message{
Headers: headers,
Queue: queue,
@@ -35,243 +29,63 @@ func NewMessage(cmd consts.CMD, payload json.RawMessage, queue string, headers m
}
}
func (m *Message) Serialize(aesKey, hmacKey []byte, encrypt bool) ([]byte, string, error) {
m.m.Lock()
defer m.m.Unlock()
buf := bytes.NewBuffer(make([]byte, 0, 512))
if err := writeLengthPrefixedJSON(buf, m.Headers); err != nil {
return nil, "", fmt.Errorf("error serializing headers: %v", err)
}
if err := writeLengthPrefixed(buf, utils.ToByte(m.Queue)); err != nil {
return nil, "", fmt.Errorf("error serializing queue: %v", err)
}
if err := binary.Write(buf, binary.LittleEndian, m.Command); err != nil {
return nil, "", fmt.Errorf("error serializing command: %v", err)
}
if err := writePayload(buf, aesKey, m.Payload, encrypt); err != nil {
return nil, "", fmt.Errorf("error serializing payload: %v", err)
}
messageBytes := buf.Bytes()
hmacSignature := CalculateHMAC(hmacKey, messageBytes)
return messageBytes, hmacSignature, nil
}
// Serialize encodes the Message to a byte slice using MessagePack.
func (m *Message) Serialize() ([]byte, error) {
m.m.RLock()
defer m.m.RUnlock()
func Deserialize(data, aesKey, hmacKey []byte, receivedHMAC string, decrypt bool) (*Message, error) {
if !VerifyHMAC(hmacKey, data, receivedHMAC) {
return nil, fmt.Errorf("HMAC verification failed")
}
buf := bytes.NewReader(data)
headers := make(map[string]string)
if err := readLengthPrefixedJSON(buf, &headers); err != nil {
return nil, fmt.Errorf("error deserializing headers: %v", err)
}
queue, err := readLengthPrefixedString(buf)
if err != nil {
return nil, fmt.Errorf("error deserializing queue: %v", err)
}
var command consts.CMD
if err := binary.Read(buf, binary.LittleEndian, &command); err != nil {
return nil, fmt.Errorf("error deserializing command: %v", err)
}
payload, err := readPayload(buf, aesKey, decrypt)
if err != nil {
return nil, fmt.Errorf("error deserializing payload: %v", err)
}
return &Message{
Headers: headers,
Queue: queue,
Command: command,
Payload: payload,
}, nil
}
func SendMessage(conn io.Writer, msg *Message, aesKey, hmacKey []byte, encrypt bool) error {
sentData, hmacSignature, err := msg.Serialize(aesKey, hmacKey, encrypt)
if err != nil {
return fmt.Errorf("error serializing message: %v", err)
}
if err := writeMessageWithHMAC(conn, sentData, hmacSignature); err != nil {
return fmt.Errorf("error sending message: %v", err)
}
return nil
}
func ReadMessage(conn io.Reader, aesKey, hmacKey []byte, decrypt bool) (*Message, error) {
data, receivedHMAC, err := readMessageWithHMAC(conn)
data, err := msgpack.Marshal(m)
if err != nil {
return nil, err
}
return Deserialize(data, aesKey, hmacKey, receivedHMAC, decrypt)
}
func writeLengthPrefixed(buf *bytes.Buffer, data []byte) error {
length := uint32(len(data))
if err := binary.Write(buf, binary.LittleEndian, length); err != nil {
return err
}
buf.Write(data)
return nil
}
func writeLengthPrefixedJSON(buf *bytes.Buffer, v any) error {
data, err := json.Marshal(v)
if err != nil {
return err
}
return writeLengthPrefixed(buf, data)
}
func readLengthPrefixed(r *bytes.Reader) ([]byte, error) {
var length uint32
if err := binary.Read(r, binary.LittleEndian, &length); err != nil {
return nil, err
}
data := make([]byte, length)
if _, err := r.Read(data); err != nil {
return nil, err
}
return data, nil
}
func readLengthPrefixedJSON(r *bytes.Reader, v any) error {
data, err := readLengthPrefixed(r)
if err != nil {
return err
}
return json.Unmarshal(data, v)
}
func readLengthPrefixedString(r *bytes.Reader) (string, error) {
data, err := readLengthPrefixed(r)
if err != nil {
return "", err
}
return utils.FromByte(data), nil
}
func writePayload(buf *bytes.Buffer, aesKey []byte, payload json.RawMessage, encrypt bool) error {
payloadBytes, err := json.Marshal(payload)
if err != nil {
return fmt.Errorf("error marshalling payload: %v", err)
}
var encryptedPayload, nonce []byte
if encrypt {
encryptedPayload, nonce, err = EncryptPayload(aesKey, payloadBytes)
if err != nil {
return err
}
} else {
encryptedPayload = payloadBytes
}
if err := writeLengthPrefixed(buf, encryptedPayload); err != nil {
return err
}
if encrypt {
buf.Write(nonce)
}
return nil
}
func readPayload(r *bytes.Reader, aesKey []byte, decrypt bool) (json.RawMessage, error) {
encryptedPayload, err := readLengthPrefixed(r)
if err != nil {
// Deserialize decodes a byte slice to a Message instance using MessagePack.
func Deserialize(data []byte) (*Message, error) {
var msg Message
if err := msgpack.Unmarshal(data, &msg); err != nil {
return nil, err
}
var payloadBytes []byte
if decrypt {
nonce := make([]byte, 12)
if _, err := r.Read(nonce); err != nil {
return nil, err
}
payloadBytes, err = DecryptPayload(aesKey, encryptedPayload, nonce)
if err != nil {
return nil, fmt.Errorf("error decrypting payload: %v", err)
}
} else {
payloadBytes = encryptedPayload
}
var payload json.RawMessage
if err := json.Unmarshal(payloadBytes, &payload); err != nil {
return nil, fmt.Errorf("error unmarshalling payload: %v", err)
}
return payload, nil
return &msg, nil
}
func writeMessageWithHMAC(conn io.Writer, messageBytes []byte, hmacSignature string) error {
if err := binary.Write(conn, binary.LittleEndian, uint32(len(messageBytes))); err != nil {
return err
}
if _, err := conn.Write(messageBytes); err != nil {
return err
}
hmacBytes, err := hex.DecodeString(hmacSignature)
// SendMessage sends a Message over a net.Conn.
func SendMessage(conn net.Conn, msg *Message) error {
data, err := msg.Serialize()
if err != nil {
return err
}
if _, err := conn.Write(hmacBytes); err != nil {
// Send the length of the data followed by the data itself
length := make([]byte, 4)
binary.BigEndian.PutUint32(length, uint32(len(data)))
if _, err := conn.Write(length); err != nil {
return err
}
if _, err := conn.Write(data); err != nil {
return err
}
return nil
}
func readMessageWithHMAC(conn io.Reader) ([]byte, string, error) {
var length uint32
if err := binary.Read(conn, binary.LittleEndian, &length); err != nil {
return nil, "", err
// ReadMessage receives a Message from a net.Conn.
func ReadMessage(conn net.Conn) (*Message, error) {
// Read the length of the incoming message
lengthBytes := make([]byte, 4)
if _, err := conn.Read(lengthBytes); err != nil {
return nil, err
}
length := binary.BigEndian.Uint32(lengthBytes)
// Read the actual message data
data := make([]byte, length)
if _, err := io.ReadFull(conn, data); err != nil {
return nil, "", err
}
hmacBytes := make([]byte, 32)
if _, err := io.ReadFull(conn, hmacBytes); err != nil {
return nil, "", err
}
receivedHMAC := hex.EncodeToString(hmacBytes)
return data, receivedHMAC, nil
}
func EncryptPayload(key []byte, plaintext []byte) ([]byte, []byte, error) {
block, err := aes.NewCipher(key)
if err != nil {
return nil, nil, err
}
aesGCM, err := cipher.NewGCM(block)
if err != nil {
return nil, nil, err
}
nonce := make([]byte, aesGCM.NonceSize())
if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
return nil, nil, err
}
ciphertext := aesGCM.Seal(nil, nonce, plaintext, nil)
return ciphertext, nonce, nil
}
func DecryptPayload(key []byte, ciphertext []byte, nonce []byte) ([]byte, error) {
block, err := aes.NewCipher(key)
if err != nil {
if _, err := conn.Read(data); err != nil {
return nil, err
}
aesGCM, err := cipher.NewGCM(block)
if err != nil {
return nil, err
}
plaintext, err := aesGCM.Open(nil, nonce, ciphertext, nil)
if err != nil {
return nil, err
}
return plaintext, nil
}
func CalculateHMAC(key []byte, data []byte) string {
h := hmac.New(sha256.New, key)
h.Write(data)
return hex.EncodeToString(h.Sum(nil))
}
func VerifyHMAC(key []byte, data []byte, receivedHMAC string) bool {
expectedHMAC := CalculateHMAC(key, data)
return hmac.Equal(utils.ToByte(receivedHMAC), utils.ToByte(expectedHMAC))
return Deserialize(data)
}