webpush-client-go/autopush/client.go

package autopush

import (
	"context"
	"crypto/ecdh"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"git.min.rip/min/webpush-client-go/rfc8291"
	"git.min.rip/min/webpush-client-go/webpush"
	"git.min.rip/min/websocket-client-go/websocket"
)

const MOZILLA_PUSH_SERVICE = "wss://push.services.mozilla.com"

type AutoPushClient struct {
	*websocket.WebSocketClient
	awaitingPong     bool
	ece              *rfc8291.RFC8291
	helloChan        chan HelloResponse
	notificationChan chan Notification
	registerChan     chan RegisterResponse
	unregisterChan   chan UnregisterResponse
}

func request[T any](c *AutoPushClient, ch chan T, timeout time.Duration, errName string, payload any) (res T, err error) {
	if err := c.SendJSON(payload); err != nil {
		log.Println("websocket send", errName, "request failed", err)
	}

	select {
	case res, ok := <-ch:
		if !ok {
			return res, fmt.Errorf("websocket recv %s failed (conn close)", errName)
		}
		return res, err
	case <-time.After(timeout * time.Second):
		return res, fmt.Errorf("%s timeout", errName)
	}
}

func unmarshaler[T any](payload json.RawMessage, label MessageType) (data *T) {
	// log.Println("autopush: before unmarshal payload", payload)
	if err := json.Unmarshal(payload, &data); err != nil {
		log.Printf("AutoPush: failed to unmarshal %s payload: %v", label, err)
	}
	return data
}

func (c *AutoPushClient) Ping() error {
	c.awaitingPong = true
	return c.SendJSON(struct{}{})
}

func (c *AutoPushClient) Hello(uaid string, channelIDs []string) (HelloResponse, error) {
	return request(c, c.helloChan, 5, "hello", HelloRequest{
		Type:       HELLO,
		UAID:       uaid,
		ChannelIDs: channelIDs,
		UseWebPush: true,
	})
}

func (c *AutoPushClient) Register(channelID string, vapidKey string) (RegisterResponse, error) {
	return request(c, c.registerChan, 5, "register", RegisterRequest{
		Type:      REGISTER,
		ChannelID: channelID,
		Key:       vapidKey,
	})
}

func (c *AutoPushClient) Unregister(channelID string) (UnregisterResponse, error) {
	return request(c, c.unregisterChan, 5, "unregister", UnregisterRequest{
		Type:      UNREGISTER,
		ChannelID: channelID,
	})
}

func (c *AutoPushClient) Decrypt(
	authSecret []byte,
	useragentPrivateKey *ecdh.PrivateKey,
	notification Notification,
) (*webpush.WebPushPayload, error) {
	data, err := base64.RawURLEncoding.DecodeString(notification.Data)
	if err != nil {
		return nil, fmt.Errorf("base64 decode error: %v", err)
	}

	plaintext, err := c.ece.Decrypt(
		data,
		rfc8291.Encoding(notification.Headers.Encoding),
		notification.Headers.Encryption,
		notification.Headers.CryptoKey,
		authSecret,
		useragentPrivateKey,
	)
	if err != nil {
		return nil, fmt.Errorf("decrypt error: %v", err)
	}

	var result webpush.WebPushPayload
	if err := json.Unmarshal(plaintext, &result); err != nil {
		return nil, fmt.Errorf("failed to unmarshal json: %v", err)
	}
	return &result, nil
}

func NewAutoPushClient(ctx context.Context) (ap *AutoPushClient, ch chan Notification) {
	ctx, cancel := context.WithCancel(ctx)
	ap = &AutoPushClient{
		ece:              rfc8291.NewRFC8291(sha256.New),
		helloChan:        make(chan HelloResponse),
		notificationChan: make(chan Notification),
		registerChan:     make(chan RegisterResponse),
		unregisterChan:   make(chan UnregisterResponse),
		WebSocketClient: websocket.NewWebSocketClient(
			ctx,
			func(ws *websocket.WebSocketClient) {
				// 3min: 10x more often than firefox
				// if this is below 45sec, autopush will disconnect us
				//
				// https://github.com/mozilla-services/autopush-rs/blob/master/autoconnect/autoconnect-ws/autoconnect-ws-sm/src/identified/on_client_msg.rs#L295
				// https://searchfox.org/firefox-main/source/dom/push/PushServiceWebSocket.sys.mjs#463
				// https://searchfox.org/firefox-main/source/modules/libpref/init/all.js#3230
				pushPing := time.NewTicker(55 * time.Second)

				go func() {
					for {
						select {
						case <-ctx.Done():
							return
						case <-pushPing.C:
							ap.Ping()
						}
					}
				}()
			},
			func(ws *websocket.WebSocketClient) {
				cancel()
				close(ap.helloChan)
				close(ap.notificationChan)
				close(ap.registerChan)
				close(ap.unregisterChan)
			},
			func(ws *websocket.WebSocketClient, payload []byte) {
				// log.Println("AutoPush Received Message:", string(payload))
				var message Message
				if err := json.Unmarshal(payload, &message); err != nil {
					log.Println("AutoPush: failed to unmarshal payload", err)
					return
				}

				if message.Type == NULL {
					message.Type = PING
				}

				switch message.Type {
				case PING:
					if ap.awaitingPong {
						ap.awaitingPong = false
					} else {
						ap.Ping()
					}

				case HELLO:
					if data := unmarshaler[HelloResponse](payload, HELLO); data != nil {
						ap.helloChan <- *data
					}

				case REGISTER:
					if data := unmarshaler[RegisterResponse](payload, REGISTER); data != nil {
						ap.registerChan <- *data
					}

				case UNREGISTER:
					if data := unmarshaler[UnregisterResponse](payload, UNREGISTER); data != nil {
						ap.unregisterChan <- *data
					}

				case NOTIFICATION:
					if data := unmarshaler[Notification](payload, NOTIFICATION); data != nil {
						ws.SendJSON(Ack{
							Type: ACK,
							Updates: []AckUpdate{
								{
									ChannelID: data.ChannelID,
									Version:   data.Version,
								},
							},
						})
						ap.notificationChan <- *data
					}

				default:
					log.Println("AutoPush: unknown data type", message.Type)
				}
			},
		),
	}

	return ap, ap.notificationChan
}