notify/envelope.go

package notify

import (
	"b612.me/notify/internal/timeutil"
	"context"
	crand "crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"time"
)

type EnvelopeKind uint8

const (
	EnvelopeSignal EnvelopeKind = iota
	EnvelopeSignalAck
	EnvelopeStreamData
	EnvelopeFileMeta
	EnvelopeFileChunk
	EnvelopeFileEnd
	EnvelopeFileAbort
	EnvelopeAck
)

type Envelope struct {
	Kind   EnvelopeKind
	ID     uint64
	Body   []byte
	Stream StreamPacket
	File   FilePacket

	controlCtx       context.Context
	controlPriority  controlPriority
	controlTimeout   time.Duration
	transportProfile *transportProtectionProfile
}

type StreamPacket struct {
	StreamID string
	Chunk    []byte
}

type FilePacket struct {
	FileID   string
	Name     string
	Size     int64
	Mode     uint32
	ModTime  int64
	Offset   int64
	Chunk    []byte
	Checksum string
	Error    string
	Stage    string
}

func wrapTransferMsgEnvelope(msg TransferMsg, enFn func(interface{}) ([]byte, error)) (Envelope, error) {
	body, err := enFn(msg)
	if err != nil {
		return Envelope{}, err
	}
	return Envelope{
		Kind:            EnvelopeSignal,
		ID:              msg.ID,
		Body:            body,
		controlPriority: controlPriorityForTransferMessage(msg),
	}, nil
}

func controlPriorityForTransferMessage(msg TransferMsg) controlPriority {
	switch msg.Type {
	case MSG_SYS, MSG_SYS_WAIT, MSG_SYS_REPLY, MSG_KEY_CHANGE, MSG_SYNC_REPLY:
		return controlPriorityCritical
	}
	if msg.Key == "heartbeat" || msg.Key == "bye" || strings.HasPrefix(msg.Key, "notify.") {
		return controlPriorityCritical
	}
	return controlPriorityNormal
}

func (env Envelope) controlContext() context.Context {
	if env.controlCtx == nil {
		return context.Background()
	}
	return env.controlCtx
}

func unwrapTransferMsgEnvelope(env Envelope, deFn func([]byte) (interface{}, error)) (TransferMsg, error) {
	if env.Kind != EnvelopeSignal {
		return TransferMsg{}, errors.New("envelope kind is not signal")
	}
	data, err := deFn(env.Body)
	if err != nil {
		return TransferMsg{}, err
	}
	msg, ok := data.(TransferMsg)
	if !ok {
		return TransferMsg{}, errors.New("invalid signal envelope payload")
	}
	return msg, nil
}

func newSignalAckEnvelope(signalID uint64) Envelope {
	return Envelope{
		Kind:            EnvelopeSignalAck,
		ID:              signalID,
		controlPriority: controlPriorityCritical,
	}
}

func newStreamDataEnvelope(streamID string, chunk []byte) Envelope {
	return Envelope{
		Kind: EnvelopeStreamData,
		Stream: StreamPacket{
			StreamID: streamID,
			Chunk:    chunk,
		},
	}
}

func newFileMetaEnvelope(fileID string, fileName string, fileSize int64, checksum string, mode uint32, modTime int64) Envelope {
	return Envelope{
		Kind: EnvelopeFileMeta,
		File: FilePacket{
			FileID:   fileID,
			Name:     filepath.Base(fileName),
			Size:     fileSize,
			Mode:     mode,
			ModTime:  modTime,
			Checksum: checksum,
		},
	}
}

func newFileChunkEnvelope(fileID string, offset int64, chunk []byte) Envelope {
	return Envelope{
		Kind: EnvelopeFileChunk,
		File: FilePacket{
			FileID: fileID,
			Offset: offset,
			Chunk:  chunk,
		},
	}
}

func newFileEndEnvelope(fileID string) Envelope {
	return Envelope{
		Kind: EnvelopeFileEnd,
		File: FilePacket{
			FileID: fileID,
		},
	}
}

func newFileAbortEnvelope(fileID string, stage string, offset int64, errMsg string) Envelope {
	return Envelope{
		Kind: EnvelopeFileAbort,
		File: FilePacket{
			FileID: fileID,
			Stage:  stage,
			Offset: offset,
			Error:  errMsg,
		},
	}
}

func newFileAckEnvelope(fileID string, stage string, offset int64, errMsg string) Envelope {
	return Envelope{
		Kind: EnvelopeAck,
		File: FilePacket{
			FileID: fileID,
			Stage:  stage,
			Offset: offset,
			Error:  errMsg,
		},
	}
}

var fileIDSerial uint64

func buildFileID(fileName string) string {
	base := fileIDBaseName(fileName)
	ts := uint64(timeutil.NowUnixNano())
	pid := uint64(os.Getpid())
	seq := atomic.AddUint64(&fileIDSerial, 1)
	rnd := uint64(randomFileIDSuffix())
	return fmt.Sprintf("%s-%x-%x-%x-%x", base, ts, pid, seq, rnd)
}

func fileIDBaseName(fileName string) string {
	base := sanitizeFileName(filepath.Base(fileName))
	switch base {
	case "", ".", "/", "\\":
		return "unnamed"
	default:
		return base
	}
}

func randomFileIDSuffix() uint32 {
	var buf [4]byte
	if _, err := crand.Read(buf[:]); err == nil {
		return binary.BigEndian.Uint32(buf[:])
	}
	seq := atomic.LoadUint64(&fileIDSerial)
	mix := uint64(timeutil.NowUnixNano()) ^ (seq << 1) ^ uint64(os.Getpid())
	return uint32(mix ^ (mix >> 32))
}