qmk-rgb/internal/via/protocol.go

package via

import (
	"fmt"

	"github.com/wobkey/rgb/internal/device"
	"github.com/wobkey/rgb/internal/hid"
)

type transport interface {
	SendReport(reportID byte, report []byte) (int, error)
	Read(buf []byte) (int, error)
	Close() error
}

// Protocol implements the VIA protocol for communication with a keyboard.
type Protocol struct {
	handle transport
	kb     device.Keyboard
}

// New creates a new VIA protocol handler for a connected device.
func New(dev device.Device) (*Protocol, error) {
	var h *hid.Device
	var err error

	if dev.Path != "" {
		h, err = hid.OpenPath(dev.Path)
	} else {
		h, err = hid.Open(dev.VendorID, dev.ProductID)
	}
	if err != nil {
		return nil, fmt.Errorf("open device: %w", err)
	}

	return &Protocol{
		handle: h,
		kb:     dev.Keyboard,
	}, nil
}

// Close releases the device handle.
func (p *Protocol) Close() error {
	if p.handle != nil {
		return p.handle.Close()
	}
	return nil
}

// Message defines VIA message types.
type Message uint8

const (
	CustomSet Message = 0x07
	CustomGet Message = 0x08
	Unhandled Message = 0xff
)

// LEDType defines QMK LED subsystem types.
type LEDType uint8

const (
	RGBLight  LEDType = 0x02
	RGBMatrix LEDType = 0x03
	SideLight LEDType = 0x04
)

// SetValue sends a set value command for the given LED type and parameter.
func (p *Protocol) SetValue(ledType LEDType, param uint8, value uint8) error {
	report := make([]byte, 32)
	report[0] = byte(CustomSet)
	report[1] = byte(ledType)
	report[2] = param
	report[3] = value

	if _, err := p.handle.SendReport(0x00, report); err != nil {
		return err
	}
	_, err := p.readResponse(CustomSet, ledType, param)
	return err
}

func (p *Protocol) DisableLighting() error {
	for _, channel := range []LEDType{RGBLight, RGBMatrix, SideLight} {
		if err := p.SetValue(channel, 0x02, 0x00); err != nil {
			return err
		}
		if err := p.SetValue(channel, 0x01, 0x00); err != nil {
			return err
		}
	}
	return nil
}

func (p *Protocol) EnableLighting() error {
	channels := []struct {
		channel LEDType
		effect  uint8
	}{
		{RGBLight, 0x04},
		{RGBMatrix, 0x05},
		{SideLight, 0x04},
	}
	for _, channel := range channels {
		if err := p.SetValue(channel.channel, 0x02, channel.effect); err != nil {
			return err
		}
		if err := p.SetValue(channel.channel, 0x01, 160); err != nil {
			return err
		}
	}
	return nil
}

func (p *Protocol) SetLightingColor(hue uint8, saturation uint8) error {
	for _, channel := range []LEDType{RGBLight, RGBMatrix, SideLight} {
		if err := p.SetColor(channel, hue, saturation); err != nil {
			return err
		}
	}
	return nil
}

func (p *Protocol) SetColor(ledType LEDType, hue uint8, saturation uint8) error {
	report := make([]byte, 32)
	report[0] = byte(CustomSet)
	report[1] = byte(ledType)
	report[2] = 0x04
	report[3] = hue
	report[4] = saturation

	if _, err := p.handle.SendReport(0x00, report); err != nil {
		return err
	}
	_, err := p.readResponse(CustomSet, ledType, 0x04)
	return err
}

// GetValue sends a get value request and reads the response.
func (p *Protocol) GetValue(ledType LEDType, param uint8) ([]byte, error) {
	report := make([]byte, 32)
	report[0] = byte(CustomGet)
	report[1] = byte(ledType)
	report[2] = param

	_, err := p.handle.SendReport(0x00, report)
	if err != nil {
		return nil, fmt.Errorf("send get request: %w", err)
	}

	buf, err := p.readResponse(CustomGet, ledType, param)
	if err != nil {
		return nil, err
	}

	valueSize := 1
	if param == 0x04 {
		valueSize = 2
	}
	return append([]byte(nil), buf[3:3+valueSize]...), nil
}

func (p *Protocol) readResponse(command Message, ledType LEDType, param uint8) ([]byte, error) {
	buf := make([]byte, 32)
	n, err := p.handle.Read(buf)
	if err != nil {
		return nil, fmt.Errorf("read response: %w", err)
	}
	if n != len(buf) {
		return nil, fmt.Errorf("short response: got %d bytes, want 32", n)
	}
	if buf[0] == byte(Unhandled) {
		return nil, fmt.Errorf("unhandled response")
	}
	if buf[0] != byte(command) {
		return nil, fmt.Errorf("unexpected response command: 0x%02x", buf[0])
	}
	if buf[1] != byte(ledType) || buf[2] != param {
		return nil, fmt.Errorf("response value mismatch: got channel 0x%02x value 0x%02x", buf[1], buf[2])
	}
	return buf, nil
}