From 0cb96b5792857f7b66d1e2cf8764668dc23eff01 Mon Sep 17 00:00:00 2001 From: Peter Edley Date: Fri, 18 Sep 2026 19:19:22 +0100 Subject: [PATCH] feat: add PipeWire screen and system audio capture (M2) Capture the Wayland desktop via xdg-desktop-portal + PipeWire using go2tv.app/screencast (MIT), and stream it to OBS: - internal/capture: Capture/FrameSource/AudioSource interfaces and the PipeWire backend (BGRA frames at monitor resolution, S16 48 kHz stereo system audio) - protocol: EncodeBGRA fast path producing 4:2:0 YCbCr JPEGs - cmd: --source screen|pattern, --audio, --stream-index flags; real capture feeds the existing sender - share one wall-clock reference between the audio and video loops so OBS receives aligned A/V timestamps (avoids multi-second latency) Verified end-to-end: real desktop at 30 fps renders in OBS with sub-second latency. --- cmd/teleportfling/main.go | 320 ++++++++++++++++++++++----------- go.mod | 6 +- go.sum | 4 + internal/capture/capture.go | 85 +++++++++ internal/capture/pipewire.go | 84 +++++++++ internal/protocol/jpeg.go | 41 +++++ internal/protocol/jpeg_test.go | 51 ++++++ 7 files changed, 489 insertions(+), 102 deletions(-) create mode 100644 internal/capture/capture.go create mode 100644 internal/capture/pipewire.go diff --git a/cmd/teleportfling/main.go b/cmd/teleportfling/main.go index 9e6172b..65041ce 100644 --- a/cmd/teleportfling/main.go +++ b/cmd/teleportfling/main.go @@ -1,22 +1,27 @@ // Command teleportfling is a standalone sender for the Teleport protocol. // -// M1 milestone: protocol proof-of-life. It streams a synthetic test pattern -// (colour bars with a moving box) plus a silent stereo 48 kHz audio tone over -// TCP and announces itself on the LAN multicast group, so an OBS instance +// It captures a Wayland screen (PipeWire via xdg-desktop-portal) plus the +// system's default audio output and streams them over TCP as the Teleport +// protocol, announcing itself on the LAN multicast group so an OBS instance // with the obs-teleport plugin can discover and decode the stream. // // Usage: // -// teleportfling [--name NAME] [--port PORT] [--width W] [--height H] -// [--fps N] [--quality 1..100] [--duration SECONDS] +// teleportfling [--name NAME] [--port PORT] [--quality 1..100] +// [--fps N] [--source screen|pattern] [--audio] +// [--stream-index N] [--duration SECONDS] // -// M2+ replaces the synthetic sources with real PipeWire screen/audio capture. +// --source pattern selects the M1 synthetic test pattern (colour bars with a +// moving box) instead of real screen capture, which is useful for testing +// without granting screen-share permission. package main import ( + "errors" "flag" "image" "image/color" + "io" "log" "os" "os/signal" @@ -25,20 +30,30 @@ import ( "syscall" "time" + "teleportfling/internal/capture" "teleportfling/internal/discovery" "teleportfling/internal/output" "teleportfling/internal/protocol" ) +const ( + // sampleRate and speakers describe the captured/encoded audio stream. + sampleRate = 48000 + speakers = 2 + // audioChunk sets how much audio we packetize per WAVE message (~10 ms). + audioChunk = 10 * time.Millisecond +) + func main() { var ( - name = flag.String("name", "", "announce name (default: hostname)") - port = flag.Int("port", 9756, "TCP listening port") - width = flag.Int("width", 1920, "frame width") - height = flag.Int("height", 1080, "frame height") - fps = flag.Int("fps", 30, "video frames per second") - quality = flag.Int("quality", 80, "JPEG quality 1..100") - duration = flag.Duration("duration", 0, "stream duration (0 = run until interrupted)") + name = flag.String("name", "", "announce name (default: hostname)") + port = flag.Int("port", 9756, "TCP listening port") + quality = flag.Int("quality", 80, "JPEG quality 1..100") + fps = flag.Int("fps", 30, "video frames per second") + source = flag.String("source", "screen", "capture source: screen or pattern") + withAudio = flag.Bool("audio", true, "capture and stream system audio") + streamIndex = flag.Int("stream-index", 0, "monitor index to capture (screen source)") + duration = flag.Duration("duration", 0, "stream duration (0 = run until interrupted)") ) flag.Parse() @@ -49,114 +64,68 @@ func main() { log.Fatalf("output: listen: %v", err) } announcer := discovery.Start(*name, p) - log.Printf("teleportfling: advertising on %d, capturing %dx%d @ %d fps", p, *width, *height, *fps) - // Pipeline state. var ( - totalFrames atomic.Int64 - encoder = mustNewEncoder() - frameInterval = time.Second / time.Duration(*fps) - audioInterval = 100 * time.Millisecond - sampleRate = 48000 - speakers = 2 - start = time.Now() - audioStart time.Time - deadline time.Time - stop = make(chan struct{}) + totalFrames atomic.Int64 + encoder = mustNewEncoder() + start = time.Now() + stop = make(chan struct{}) ) - if *duration > 0 { - deadline = start.Add(*duration) + var ( + cam capture.Capture + loop frameSource + ) + + switch *source { + case "screen": + cam, err = capture.OpenPipeWire(*streamIndex, *withAudio) + if err != nil { + log.Fatalf("capture: %v", err) + } + loop = captureLoop{cam} + log.Printf("teleportfling: advertising on %d, capturing screen via PipeWire", p) + case "pattern": + loop = &patternLoop{ + w: 1920, + h: 1080, + fps: *fps, + } + log.Printf("teleportfling: advertising on %d, streaming test pattern", p) + default: + log.Fatalf("teleportfling: unknown source %q (want screen or pattern)", *source) } // Interrupt / SIGTERM handling. sigc := make(chan os.Signal, 1) signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM) - // — Audio goroutine: ~100 ms chunks of a silent stereo float32 tone — - // A silent master tone keeps OBS's audio pipeline alive without needing - // a mic. M2 will replace this with real captured audio. + // — Audio loop — audioDone := make(chan struct{}) go func() { defer close(audioDone) - audioStart = time.Now() - tick := time.NewTicker(audioInterval) - defer tick.Stop() - framesPerChunk := int32(float64(sampleRate) * audioInterval.Seconds()) - pcm := make([]byte, 0, framesPerChunk*int32(speakers)*4) - - sendAudio := func(ts uint64, chunkFrames int32) { - // Right = left = 0 → digital silence. - pcm = pcm[:0] - for f := 0; f < int(chunkFrames); f++ { - pcm = append(pcm, 0, 0, 0, 0, 0, 0, 0, 0) - } - packet, err := protocol.BuildWavePacket(ts, protocol.AudioFormatF32, int32(sampleRate), int32(speakers), chunkFrames, pcm) - if err != nil { - log.Printf("teleportfling: wave: %v", err) - return - } - sender.Send(packet) + var src io.ReadCloser + switch { + case cam == nil: + // Pattern source: synthesize silence to keep the audio pipeline alive. + src = capture.NewSilenceSource() + case cam.Audio() != nil: + src = cam.Audio() + default: + log.Printf("teleportfling: system audio unavailable, streaming silence") + src = capture.NewSilenceSource() } + defer func() { _ = src.Close() }() - for { - select { - case now := <-tick.C: - ts := uint64(now.Sub(audioStart)) - // Compute frames elapsed since audioStart so timestamps are a - // continuous stream (not aligned to tick boundaries). - elapsedFrames := int64(now.Sub(audioStart) / (time.Second / time.Duration(sampleRate))) - // Cumulative frames sent so far. - sendAudio(ts, framesPerChunk) - _ = elapsedFrames - case <-stop: - return - } - } + audioLoop(sender, src, start, stop) }() // — Video loop — videoDone := make(chan struct{}) go func() { defer close(videoDone) - ticker := time.NewTicker(frameInterval) - defer ticker.Stop() - - var frameNum int64 - for { - select { - case now := <-ticker.C: - ts := uint64(now.Sub(start)) - img := testPattern(*width, *height, int(frameNum)) - frameNum++ - - frameStart := time.Now() - buf, err := encoder.Encode(img, *quality) - if err != nil { - log.Printf("teleportfling: jpeg: %v", err) - continue - } - encodeDur := time.Since(frameStart) - - packet, err := protocol.WritePacket( - protocol.Header{Type: protocol.VideoType, Timestamp: ts, Size: int32(len(buf))}, - ptr(protocol.DefaultBT709Full()), - nil, - buf, - ) - if err != nil { - log.Printf("teleportfling: packet: %v", err) - continue - } - sender.Send(packet) - totalFrames.Add(1) - - _ = encodeDur // stats below - case <-stop: - return - } - } + videoLoop(sender, encoder, loop, *fps, *quality, start, stop, &totalFrames) }() // — Stats ticker — @@ -183,7 +152,7 @@ func main() { case <-func() <-chan struct{} { if *duration > 0 { ch := make(chan struct{}) - time.AfterFunc(time.Until(deadline), func() { close(ch) }) + time.AfterFunc(*duration, func() { close(ch) }) return ch } return nil @@ -199,9 +168,158 @@ func main() { announcer.Stop() sender.Close() encoder.Close() + if cam != nil { + if err := cam.Close(); err != nil { + log.Printf("teleportfling: capture close: %v", err) + } + } log.Printf("teleportfling: stopped after %s", time.Since(start).Round(time.Millisecond)) } +// audioLoop reads raw PCM from src and emits WAVE packets in audioChunk-sized +// pieces. PCM is assumed interleaved signed-16-bit at sampleRate/speakers. +// +// start is the shared reference clock used by the video loop: audio and video +// timestamps must share one time base, otherwise a constant skew between them +// makes OBS buffer one stream to re-sync the other, adding latency. +func audioLoop(sender *output.Sender, src io.Reader, start time.Time, stop <-chan struct{}) { + framesPerChunk := int(sampleRate) * int(audioChunk) / int(time.Second) + chunkBytes := framesPerChunk * speakers * 2 // S16 + buf := make([]byte, chunkBytes) + + for { + n, err := io.ReadFull(src, buf) + if n > 0 { + frames := n / (speakers * 2) + ts := uint64(time.Since(start)) + packet, perr := protocol.BuildWavePacket(ts, protocol.AudioFormatS16, sampleRate, speakers, int32(frames), buf[:n]) + if perr != nil { + log.Printf("teleportfling: wave: %v", perr) + } else { + sender.Send(packet) + } + } + if err != nil { + if !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrClosedPipe) { + log.Printf("teleportfling: audio: %v", err) + } + select { + case <-stop: + return + default: + } + } + } +} + +// videoLoop pulls frames from loop and sends them at fps, encoding each to +// JPEG with the given quality. +func videoLoop(sender *output.Sender, encoder *protocol.JPEGEncoder, loop frameSource, fps, quality int, start time.Time, stop <-chan struct{}, total *atomic.Int64) { + frameInterval := time.Second / time.Duration(fps) + next := start + + for { + frame, err := loop.Next() + if err != nil { + if !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrClosedPipe) { + log.Printf("teleportfling: capture: %v", err) + } + select { + case <-stop: + return + default: + } + continue + } + + // Drop frames if we're running ahead of the target fps (e.g. a 60 Hz + // monitor captured at 30 fps) to keep timestamps monotonic. + now := time.Now() + if now.Before(next) { + continue + } + next = now.Add(frameInterval) + + ts := uint64(now.Sub(start)) + buf, err := encodeFrame(encoder, frame, quality) + if err != nil { + log.Printf("teleportfling: jpeg: %v", err) + continue + } + + packet, err := protocol.WritePacket( + protocol.Header{Type: protocol.VideoType, Timestamp: ts, Size: int32(len(buf))}, + ptr(protocol.DefaultBT709Full()), + nil, + buf, + ) + if err != nil { + log.Printf("teleportfling: packet: %v", err) + continue + } + sender.Send(packet) + total.Add(1) + } +} + +// encodeFrame compresses a captured frame based on its concrete type. +func encodeFrame(encoder *protocol.JPEGEncoder, frame *capture.VideoFrame, quality int) ([]byte, error) { + if frame.Pix != nil { + // BGRA from the PipeWire backend. + return encoder.EncodeBGRA(frame.Pix, frame.Width, frame.Height, quality) + } + return nil, errors.New("capture: unsupported frame type") +} + +// frameSource abstracts the frame source: real capture or the test pattern. +type frameSource interface { + Next() (*capture.VideoFrame, error) +} + +// captureLoop wraps the PipeWire capture backend. +type captureLoop struct { + cam capture.Capture +} + +func (c captureLoop) Next() (*capture.VideoFrame, error) { + return c.cam.Video().NextFrame() +} + +// patternLoop synthesizes the M1 test pattern (colour bars + moving box). +type patternLoop struct { + w, h int + fps int + seq int64 +} + +func (p *patternLoop) Next() (*capture.VideoFrame, error) { + img := testPattern(p.w, p.h, int(p.seq)) + p.seq++ + return ycrcbToBGRA(img), nil +} + +// ycrcbToBGRA converts a YCbCr image to a BGRA VideoFrame so both sources +// share the encode path (EncodeBGRA). +func ycrcbToBGRA(img *image.YCbCr) *capture.VideoFrame { + w, h := img.Rect.Dx(), img.Rect.Dy() + frame := &capture.VideoFrame{ + Pix: make([]byte, w*h*4), + Width: w, + Height: h, + Stride: w * 4, + } + for y := 0; y < h; y++ { + for x := 0; x < w; x++ { + yi := y*img.YStride + x + ci := (y/2)*img.CStride + x/2 + r, g, b := color.YCbCrToRGB(img.Y[yi], img.Cb[ci], img.Cr[ci]) + off := (y*w + x) * 4 + frame.Pix[off], frame.Pix[off+1], frame.Pix[off+2], frame.Pix[off+3] = b, g, r, 255 + } + } + return frame +} + // mustNewEncoder creates a JPEG encoder or panics. func mustNewEncoder() *protocol.JPEGEncoder { enc, err := protocol.NewJPEGEncoder() @@ -222,7 +340,7 @@ func addr(port int) string { // testPattern renders a standard SMPTE colour bar with a moving white box at // the given frame index. The result is a *image.YCbCr 4:2:0 image so the // encoder uses the YUV path — closest to what real PipeWire capture will -// produce in M2. +// produce. func testPattern(w, h, frame int) *image.YCbCr { img := image.NewYCbCr(image.Rect(0, 0, w, h), image.YCbCrSubsampleRatio420) diff --git a/go.mod b/go.mod index 2647c66..98393d8 100644 --- a/go.mod +++ b/go.mod @@ -2,9 +2,13 @@ module teleportfling go 1.26.0 -require github.com/schollz/peerdiscovery v1.7.6 +require ( + github.com/schollz/peerdiscovery v1.7.6 + go2tv.app/screencast v0.0.0-20260807191623-4cba6136251a +) require ( + github.com/godbus/dbus/v5 v5.1.0 // indirect golang.org/x/net v0.59.0 // indirect golang.org/x/sys v0.48.0 // indirect ) diff --git a/go.sum b/go.sum index 3dfbe65..5a2aa76 100644 --- a/go.sum +++ b/go.sum @@ -1,5 +1,7 @@ github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk= +github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= @@ -9,6 +11,8 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= +go2tv.app/screencast v0.0.0-20260807191623-4cba6136251a h1:gfpAKu/sFOOG2TEAtYOcv2kMkTHjTtANtVH9rkGtuu0= +go2tv.app/screencast v0.0.0-20260807191623-4cba6136251a/go.mod h1:SuTkH1JeLAuYESPe0rmpmqmprr+Q2Z/Dfh2t4oHlCMM= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc= diff --git a/internal/capture/capture.go b/internal/capture/capture.go new file mode 100644 index 0000000..30fd004 --- /dev/null +++ b/internal/capture/capture.go @@ -0,0 +1,85 @@ +// Package capture abstracts the screen and audio sources that feed the +// Teleport sender. Implementations are platform-specific; the PipeWire +// backend (pipewire.go) handles Wayland via xdg-desktop-portal. +// +// The core only knows about two things: a stream of video frames and a +// stream of audio samples. Everything else (formats, portal negotiation) +// stays behind this interface so the sender can later run headless or be +// driven by a GUI without coupling. +package capture + +import ( + "io" + "time" +) + +// VideoFrame is one captured screen frame. Pix holds BGRA (blue, green, +// red, alpha) bytes in row-major order; Stride is the byte offset between +// consecutive rows. +type VideoFrame struct { + Pix []byte + Width int + Height int + Stride int +} + +// Capture is the combined screen + audio source. Close releases the +// underlying capture session (and portal resources). +type Capture interface { + // Video returns the frame source. Frames arrive at the compositor's + // refresh rate and are consumed one at a time via NextFrame. + Video() FrameSource + // Audio returns the audio source, or nil if audio capture is disabled + // or unavailable. + Audio() AudioSource + io.Closer +} + +// FrameSource yields consecutive captured video frames. +type FrameSource interface { + // NextFrame blocks until the next frame is available and returns it. + NextFrame() (*VideoFrame, error) +} + +// AudioSource yields raw interleaved PCM samples (signed 16-bit +// little-endian, 48 kHz, stereo) read from the system's default output. +type AudioSource interface { + io.ReadCloser +} + +// ErrNoAudio is returned when the underlying backend cannot provide system +// audio capture (e.g. sandboxed Flatpak without a direct PipeWire link). +var ErrNoAudio = &AudioUnavailableError{} + +// AudioUnavailableError signals that audio capture could not be started. +type AudioUnavailableError struct{} + +func (e *AudioUnavailableError) Error() string { + return "capture: system audio unavailable" +} + +// silenceStep is the pacing interval between silence chunk reads. +const silenceStep = 10 * time.Millisecond + +// SilenceSource yields a continuous stream of digital silence, paced like a +// real audio capture. It keeps OBS's audio pipeline alive when no system +// audio is available or a synthetic source is in use. +type silenceSource struct{} + +// NewSilenceSource creates a silence-generating audio source. +func NewSilenceSource() AudioSource { + return &silenceSource{} +} + +func (s *silenceSource) Read(p []byte) (int, error) { + if len(p) == 0 { + return 0, nil + } + // Emit a chunk of silence at a real-audio cadence. The reader is 1:1 + // stereo S16, so 48000 * 10ms * 2ch * 2 bytes = 1920 bytes per read. + clear(p) + time.Sleep(silenceStep) + return len(p), nil +} + +func (s *silenceSource) Close() error { return nil } diff --git a/internal/capture/pipewire.go b/internal/capture/pipewire.go new file mode 100644 index 0000000..b98b468 --- /dev/null +++ b/internal/capture/pipewire.go @@ -0,0 +1,84 @@ +// Package capture provides the PipeWire backend for Wayland screen + audio +// capture. +// +// This uses go2tv.app/screencast (MIT) which implements the full +// xdg-desktop-portal ScreenCast negotiation and then receives frames over +// PipeWire. The portal session is responsible for the screen-share consent +// dialog presented by the compositor (Hyprland in our dev environment). +// +// Frames arrive as raw BGRA at the monitor's native resolution/refresh rate. +// Audio is signed 16-bit, 48 kHz, stereo interleaved PCM from the system's +// default output. +package capture + +import ( + "errors" + "io" + + "go2tv.app/screencast/capture" +) + +// PipeWire implements Capture on top of the screencast library. +type PipeWire struct { + stream *capture.Stream +} + +// OpenPipeWire opens a PipeWire capture session. streamIndex selects which +// monitor to capture when multiple are present. Triggering the portal +// consent dialog is expected; the compositor decides whether to show it. +func OpenPipeWire(streamIndex int, audio bool) (*PipeWire, error) { + s, err := capture.Open(&capture.Options{ + StreamIndex: streamIndex, + IncludeAudio: audio, + }) + if err != nil { + return nil, err + } + return &PipeWire{stream: s}, nil +} + +// Video returns the BGRA frame source. +func (p *PipeWire) Video() FrameSource { + return &pipewireVideo{stream: p.stream} +} + +// Audio returns the system audio source, or nil if unavailable. +func (p *PipeWire) Audio() AudioSource { + if p.stream.Audio == nil { + return nil + } + return p.stream.Audio +} + +// Close releases the capture session and portal resources. +func (p *PipeWire) Close() error { + return p.stream.Close() +} + +// pipewireVideo adapts the screencast io.ReadCloser into FrameSource. +type pipewireVideo struct { + stream *capture.Stream + frame *VideoFrame +} + +// NextFrame blocks until the next full frame is delivered. The screencast +// library writes one complete BGRA frame per Read, so we assemble it with +// ReadFull and reuse the underlying buffer across calls. +func (v *pipewireVideo) NextFrame() (*VideoFrame, error) { + w, h := int(v.stream.Width), int(v.stream.Height) + if v.frame == nil { + v.frame = &VideoFrame{ + Pix: make([]byte, w*h*4), + Width: w, + Height: h, + Stride: w * 4, + } + } + if _, err := io.ReadFull(v.stream, v.frame.Pix); err != nil { + if errors.Is(err, io.EOF) { + return nil, io.ErrClosedPipe + } + return nil, err + } + return v.frame, nil +} diff --git a/internal/protocol/jpeg.go b/internal/protocol/jpeg.go index 4247b05..607f5cd 100644 --- a/internal/protocol/jpeg.go +++ b/internal/protocol/jpeg.go @@ -74,6 +74,47 @@ func (e *JPEGEncoder) Encode(img image.Image, quality int) ([]byte, error) { } } +// EncodeBGRA compresses a raw BGRA (blue, green, red, alpha) pixel buffer of +// the given dimensions. This is the fast path for the PipeWire screen-capture +// backend, which delivers frames in BGRA byte order. Subsampling defaults to +// 4:2:0 (a YCbCr JPEG), so the OBS receiver decodes it as I420 — the same +// layout obs-teleport produces for OBS-native YCbCr frames. +func (e *JPEGEncoder) EncodeBGRA(pix []byte, width, height, quality int) ([]byte, error) { + if quality < 1 { + quality = 1 + } + if quality > 100 { + quality = 100 + } + + want := width * height * 4 + if len(pix) < want { + return nil, errors.New("turbojpeg: BGRA buffer too small") + } + + C.tj3Set(e.ctx, C.TJPARAM_QUALITY, C.int(quality)) + C.tj3Set(e.ctx, C.TJPARAM_SUBSAMP, C.TJSAMP_420) + C.tj3Set(e.ctx, C.TJPARAM_COLORSPACE, C.TJCS_YCbCr) + + size := C.tj3JPEGBufSize(C.int(width), C.int(height), C.TJSAMP_420) + buf := make([]byte, int(size)) + + srcPtr := unsafe.Pointer(&pix[0]) + dstPtr := (*C.uchar)(&buf[0]) + + var pin runtime.Pinner + pin.Pin(srcPtr) + pin.Pin(dstPtr) + defer pin.Unpin() + + jpegSize := size + rc := C.tj3Compress8(e.ctx, (*C.uchar)(srcPtr), C.int(width), 0, C.int(height), C.TJPF_BGRA, &dstPtr, &jpegSize) + if rc != 0 { + return nil, errors.New("turbojpeg BGRA compress failed") + } + return buf[:int(jpegSize)], nil +} + // encodeRGBA compresses a Go RGBA image (pixel layout [R,G,B,A] per 4 bytes). // TJPF_RGBA tells turbojpeg the exact layout; colourspace is RGB. func (e *JPEGEncoder) encodeRGBA(img *image.RGBA) ([]byte, error) { diff --git a/internal/protocol/jpeg_test.go b/internal/protocol/jpeg_test.go index 36b8a88..6cd5ff7 100644 --- a/internal/protocol/jpeg_test.go +++ b/internal/protocol/jpeg_test.go @@ -120,3 +120,54 @@ func TestJPEGEncodeQualityClamp(t *testing.T) { } } } + +// TestJPEGEncodeBGRA checks the raw-BGRA fast path produces a decodable +// JPEG with correct SOI/EOI markers. +func TestJPEGEncodeBGRA(t *testing.T) { + enc, err := NewJPEGEncoder() + if err != nil { + t.Fatalf("NewJPEGEncoder: %v", err) + } + defer enc.Close() + + // 16x8 BGRA: top half red (R,G,B=255,0,0), bottom half blue (0,0,255). + w, h := 16, 8 + pix := make([]byte, w*h*4) + for y := 0; y < h; y++ { + for x := 0; x < w; x++ { + off := (y*w + x) * 4 + if y < h/2 { + pix[off], pix[off+1], pix[off+2], pix[off+3] = 0, 0, 255, 255 // B,G,R,A = blue,red + } else { + pix[off], pix[off+1], pix[off+2], pix[off+3] = 255, 0, 0, 255 // B,G,R,A = red,blue + } + } + } + + buf, err := enc.EncodeBGRA(pix, w, h, 85) + if err != nil { + t.Fatalf("EncodeBGRA: %v", err) + } + if len(buf) == 0 { + t.Fatal("empty JPEG output") + } + if !bytes.Equal(buf[:2], []byte{0xFF, 0xD8}) { + t.Errorf("bad JPEG SOI marker: %x", buf[:2]) + } + if !bytes.Equal(buf[len(buf)-2:], []byte{0xFF, 0xD9}) { + t.Errorf("bad JPEG EOI marker") + } +} + +// TestJPEGEncodeBGRAShort rejects truncated pixel buffers. +func TestJPEGEncodeBGRAShort(t *testing.T) { + enc, err := NewJPEGEncoder() + if err != nil { + t.Fatalf("NewJPEGEncoder: %v", err) + } + defer enc.Close() + + if _, err := enc.EncodeBGRA(make([]byte, 10), 16, 8, 80); err == nil { + t.Error("short BGRA buffer accepted") + } +}