mirror of
https://github.com/oarkflow/mq.git
synced 2025-10-06 16:36:53 +08:00
feat: update
This commit is contained in:
274
codec/codec.go
274
codec/codec.go
@@ -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)
|
||||
}
|
||||
|
Reference in New Issue
Block a user