signalmeow/attachments: add basic support for TUS uploads

mautrix · Oct 18, 2024 · 58dacbd · 58dacbd
1 parent 96221b6
commit 58dacbd
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 33 deletions.
diff --git a/pkg/signalmeow/attachments.go b/pkg/signalmeow/attachments.go
@@ -23,6 +23,7 @@ import (
 	"crypto/cipher"
 	"crypto/hmac"
 	"crypto/sha256"
+	"encoding/base64"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -105,7 +106,7 @@ func decryptAttachment(body, key, digest []byte, size uint32) ([]byte, error) {
 	return decrypted[:size], nil
 }
 
-type attachmentV3UploadAttributes struct {
+type attachmentV4UploadAttributes struct {
 	Cdn                  uint32            `json:"cdn"`
 	Key                  string            `json:"key"`
 	Headers              map[string]string `json:"headers"`
@@ -154,31 +155,61 @@ func (cli *Client) UploadAttachment(ctx context.Context, body []byte) (*signalpb
 	opts := &web.HTTPReqOpt{Username: &username, Password: &password}
 	resp, err := web.SendHTTPRequest(ctx, http.MethodGet, attributesPath, opts)
 	if err != nil {
-		log.Err(err).Msg("Error sending request fetching upload attributes")
-		return nil, err
+		log.Err(err).Msg("Failed to request upload attributes")
+		return nil, fmt.Errorf("failed to request upload attributes: %w", err)
 	}
-	var uploadAttributes attachmentV3UploadAttributes
+	var uploadAttributes attachmentV4UploadAttributes
 	err = web.DecodeHTTPResponseBody(ctx, &uploadAttributes, resp)
 	if err != nil {
-		log.Err(err).Msg("Error decoding response body fetching upload attributes")
+		log.Err(err).Msg("Failed to decode upload attributes")
+		return nil, fmt.Errorf("failed to decode upload attributes: %w", err)
+	}
+	if uploadAttributes.Cdn == 3 {
+		log.Trace().Msg("Using TUS upload")
+		err = cli.uploadAttachmentTUS(ctx, uploadAttributes, encryptedWithMAC, username, password)
+	} else {
+		log.Trace().Msg("Using legacy upload")
+		err = cli.uploadAttachmentLegacy(ctx, uploadAttributes, encryptedWithMAC, username, password)
+	}
+	if err != nil {
+		log.Err(err).Msg("Failed to upload attachment")
 		return nil, err
 	}
 
+	digest := sha256.Sum256(encryptedWithMAC)
+
+	attachmentPointer := &signalpb.AttachmentPointer{
+		AttachmentIdentifier: &signalpb.AttachmentPointer_CdnKey{
+			CdnKey: uploadAttributes.Key,
+		},
+		Key:       keys,
+		Digest:    digest[:],
+		Size:      &plaintextLength,
+		CdnNumber: &uploadAttributes.Cdn,
+	}
+
+	return attachmentPointer, nil
+}
+
+func (cli *Client) uploadAttachmentLegacy(
+	ctx context.Context,
+	uploadAttributes attachmentV4UploadAttributes,
+	encryptedWithMAC []byte,
+	username string,
+	password string,
+) error {
 	// Allocate attachment on CDN
-	resp, err = web.SendHTTPRequest(ctx, http.MethodPost, "", &web.HTTPReqOpt{
+	resp, err := web.SendHTTPRequest(ctx, http.MethodPost, "", &web.HTTPReqOpt{
 		OverrideURL: uploadAttributes.SignedUploadLocation,
 		ContentType: web.ContentTypeOctetStream,
 		Headers:     uploadAttributes.Headers,
 		Username:    &username,
 		Password:    &password,
 	})
 	if err != nil {
-		log.Err(err).Msg("Error sending request allocating attachment")
-		return nil, err
-	}
-	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		log.Error().Int("status_code", resp.StatusCode).Msg("Error allocating attachment")
-		return nil, fmt.Errorf("error allocating attachment: %s", resp.Status)
+		return fmt.Errorf("failed to send allocate request: %w", err)
+	} else if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return fmt.Errorf("allocate request returned HTTP %d", resp.StatusCode)
 	}
 
 	// Upload attachment to CDN
@@ -190,27 +221,39 @@ func (cli *Client) UploadAttachment(ctx context.Context, body []byte) (*signalpb
 		Password:    &password,
 	})
 	if err != nil {
-		log.Err(err).Msg("Error sending request uploading attachment")
-		return nil, err
+		return fmt.Errorf("failed to send upload request: %w", err)
+	} else if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return fmt.Errorf("upload request returned HTTP %d", resp.StatusCode)
 	}
-	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
-		log.Error().Int("status_code", resp.StatusCode).Msg("Error uploading attachment")
-		return nil, fmt.Errorf("error uploading attachment: %s", resp.Status)
-	}
-
-	digest := sha256.Sum256(encryptedWithMAC)
+	return nil
+}
 
-	attachmentPointer := &signalpb.AttachmentPointer{
-		AttachmentIdentifier: &signalpb.AttachmentPointer_CdnKey{
-			CdnKey: uploadAttributes.Key,
-		},
-		Key:       keys,
-		Digest:    digest[:],
-		Size:      &plaintextLength,
-		CdnNumber: &uploadAttributes.Cdn,
+func (cli *Client) uploadAttachmentTUS(
+	ctx context.Context,
+	uploadAttributes attachmentV4UploadAttributes,
+	encryptedWithMAC []byte,
+	username string,
+	password string,
+) error {
+	uploadAttributes.Headers["Tus-Resumable"] = "1.0.0"
+	uploadAttributes.Headers["Upload-Length"] = fmt.Sprintf("%d", len(encryptedWithMAC))
+	uploadAttributes.Headers["Upload-Metadata"] = "filename " + base64.StdEncoding.EncodeToString([]byte(uploadAttributes.Key))
+
+	resp, err := web.SendHTTPRequest(ctx, http.MethodPost, "", &web.HTTPReqOpt{
+		OverrideURL: uploadAttributes.SignedUploadLocation,
+		Body:        encryptedWithMAC,
+		ContentType: web.ContentTypeOffsetOctetStream,
+		Headers:     uploadAttributes.Headers,
+		Username:    &username,
+		Password:    &password,
+	})
+	// TODO actually support resuming on error
+	if err != nil {
+		return fmt.Errorf("failed to send upload request: %w", err)
+	} else if resp.StatusCode < 200 || resp.StatusCode >= 300 {
+		return fmt.Errorf("upload request returned HTTP %d", resp.StatusCode)
 	}
-
-	return attachmentPointer, nil
+	return nil
 }
 
 func (cli *Client) UploadGroupAvatar(ctx context.Context, avatarBytes []byte, gid types.GroupIdentifier) (*string, error) {

diff --git a/pkg/signalmeow/web/web.go b/pkg/signalmeow/web/web.go
@@ -95,9 +95,10 @@ func init() {
 type ContentType string
 
 const (
-	ContentTypeJSON        ContentType = "application/json"
-	ContentTypeProtobuf    ContentType = "application/x-protobuf"
-	ContentTypeOctetStream ContentType = "application/octet-stream"
+	ContentTypeJSON              ContentType = "application/json"
+	ContentTypeProtobuf          ContentType = "application/x-protobuf"
+	ContentTypeOctetStream       ContentType = "application/octet-stream"
+	ContentTypeOffsetOctetStream ContentType = "application/offset+octet-stream"
 )
 
 type HTTPReqOpt struct {