From 72ae9383026a94d9e6d1b22c727269dc9c8e3cdc Mon Sep 17 00:00:00 2001 From: Ben Visness Date: Mon, 23 Aug 2021 16:52:57 -0500 Subject: [PATCH] Save Discord attachments and embeds --- src/assets/assets.go | 12 +- src/discord/gateway.go | 20 +-- src/discord/library.go | 2 +- src/discord/payloads.go | 268 ++++++++++++++++++++++++++++++++++++++-- src/discord/showcase.go | 181 ++++++++++++++++++++++++--- 5 files changed, 436 insertions(+), 47 deletions(-) diff --git a/src/assets/assets.go b/src/assets/assets.go index cf2bc83..bf02a50 100644 --- a/src/assets/assets.go +++ b/src/assets/assets.go @@ -48,9 +48,9 @@ func init() { } type CreateInput struct { - Content []byte - Filename string - MimeType string + Content []byte + Filename string + ContentType string // Optional params UploaderID *int // HMN user id @@ -78,8 +78,8 @@ func Create(ctx context.Context, dbConn db.ConnOrTx, in CreateInput) (*models.As if len(in.Content) == 0 { return nil, InvalidAssetError(fmt.Errorf("could not upload asset '%s': no bytes of data were provided", filename)) } - if in.MimeType == "" { - return nil, InvalidAssetError(fmt.Errorf("could not upload asset '%s': no mime type provided", filename)) + if in.ContentType == "" { + return nil, InvalidAssetError(fmt.Errorf("could not upload asset '%s': no content type provided", filename)) } // Upload the asset to the DO space @@ -128,7 +128,7 @@ func Create(ctx context.Context, dbConn db.ConnOrTx, in CreateInput) (*models.As key, filename, len(in.Content), - in.MimeType, + in.ContentType, checksum, in.Width, in.Height, diff --git a/src/discord/gateway.go b/src/discord/gateway.go index 5962086..abf9185 100644 --- a/src/discord/gateway.go +++ b/src/discord/gateway.go @@ -93,7 +93,7 @@ func RunDiscordBot(ctx context.Context, dbConn *pgxpool.Pool) <-chan struct{} { var outgoingMessagesReady = make(chan struct{}, 1) -type discordBotInstance struct { +type botInstance struct { conn *websocket.Conn dbConn *pgxpool.Pool @@ -116,8 +116,8 @@ type discordBotInstance struct { wg sync.WaitGroup } -func newBotInstance(dbConn *pgxpool.Pool) *discordBotInstance { - return &discordBotInstance{ +func newBotInstance(dbConn *pgxpool.Pool) *botInstance { + return &botInstance{ dbConn: dbConn, forceHeartbeat: make(chan struct{}), didAckHeartbeat: true, @@ -129,7 +129,7 @@ Runs a bot instance to completion. It will start up a gateway connection and ret connection is closed. It only returns an error when something unexpected occurs; if so, you should do exponential backoff before reconnecting. Otherwise you can reconnect right away. */ -func (bot *discordBotInstance) Run(ctx context.Context) (err error) { +func (bot *botInstance) Run(ctx context.Context) (err error) { defer utils.RecoverPanicAsError(&err) ctx, bot.cancel = context.WithCancel(ctx) @@ -223,7 +223,7 @@ and RESUMED messages in our main message receiving loop instead of here. That way, we could receive exactly one message after sending Resume, either a Resume ACK or an Invalid Session, and from there it would be crystal clear what to do. Alas!) */ -func (bot *discordBotInstance) connect(ctx context.Context) error { +func (bot *botInstance) connect(ctx context.Context) error { res, err := GetGatewayBot(ctx) if err != nil { return oops.New(err, "failed to get gateway URL") @@ -328,7 +328,7 @@ func (bot *discordBotInstance) connect(ctx context.Context) error { Sends outgoing gateway messages and channel messages. Handles heartbeats. This function should be run as its own goroutine. */ -func (bot *discordBotInstance) doSender(ctx context.Context) { +func (bot *botInstance) doSender(ctx context.Context) { defer bot.wg.Done() defer bot.cancel() @@ -507,7 +507,7 @@ func (bot *discordBotInstance) doSender(ctx context.Context) { } } -func (bot *discordBotInstance) receiveGatewayMessage(ctx context.Context) (*GatewayMessage, error) { +func (bot *botInstance) receiveGatewayMessage(ctx context.Context) (*GatewayMessage, error) { _, msgBytes, err := bot.conn.ReadMessage() if err != nil { return nil, err @@ -524,7 +524,7 @@ func (bot *discordBotInstance) receiveGatewayMessage(ctx context.Context) (*Gate return &msg, nil } -func (bot *discordBotInstance) sendGatewayMessage(ctx context.Context, msg GatewayMessage) error { +func (bot *botInstance) sendGatewayMessage(ctx context.Context, msg GatewayMessage) error { logging.ExtractLogger(ctx).Debug().Interface("msg", msg).Msg("sending gateway message") return bot.conn.WriteMessage(websocket.TextMessage, msg.ToJSON()) } @@ -534,7 +534,7 @@ Processes a single event message from Discord. If this returns an error, it mean really gone wrong, bad enough that the connection should be shut down. Otherwise it will just log any errors that occur. */ -func (bot *discordBotInstance) processEventMsg(ctx context.Context, msg *GatewayMessage) error { +func (bot *botInstance) processEventMsg(ctx context.Context, msg *GatewayMessage) error { if msg.Opcode != OpcodeDispatch { panic(fmt.Sprintf("processEventMsg must only be used on Dispatch messages (opcode %d). Validate this before you call this function.", OpcodeDispatch)) } @@ -562,7 +562,7 @@ func (bot *discordBotInstance) processEventMsg(ctx context.Context, msg *Gateway return nil } -func (bot *discordBotInstance) messageCreateOrUpdate(ctx context.Context, msg *Message) error { +func (bot *botInstance) messageCreateOrUpdate(ctx context.Context, msg *Message) error { if msg.OriginalHasFields("author") && msg.Author.ID == config.Config.Discord.BotUserID { // Don't process your own messages return nil diff --git a/src/discord/library.go b/src/discord/library.go index 3eaa30c..bff896c 100644 --- a/src/discord/library.go +++ b/src/discord/library.go @@ -6,7 +6,7 @@ import ( "git.handmade.network/hmn/hmn/src/oops" ) -func (bot *discordBotInstance) processLibraryMsg(ctx context.Context, msg *Message) error { +func (bot *botInstance) processLibraryMsg(ctx context.Context, msg *Message) error { switch msg.Type { case MessageTypeDefault, MessageTypeReply, MessageTypeApplicationCommand: default: diff --git a/src/discord/payloads.go b/src/discord/payloads.go index 18b1b4c..765608e 100644 --- a/src/discord/payloads.go +++ b/src/discord/payloads.go @@ -112,8 +112,9 @@ type Resume struct { type ChannelType int +// https://discord.com/developers/docs/resources/channel#channel-object-channel-types const ( - ChannelTypeGuildext ChannelType = 0 + ChannelTypeGuildText ChannelType = 0 ChannelTypeDM ChannelType = 1 ChannelTypeGuildVoice ChannelType = 2 ChannelTypeGroupDM ChannelType = 3 @@ -126,6 +127,7 @@ const ( ChannelTypeGuildStageVoice ChannelType = 13 ) +// https://discord.com/developers/docs/resources/channel#channel-object type Channel struct { ID string `json:"id"` Type ChannelType `json:"type"` @@ -138,6 +140,7 @@ type Channel struct { type MessageType int +// https://discord.com/developers/docs/resources/channel#message-object-message-types const ( MessageTypeDefault MessageType = 0 @@ -181,7 +184,7 @@ type Message struct { Type MessageType `json:"type"` Attachments []Attachment `json:"attachments"` - // TODO: Embeds + Embeds []Embed `json:"embeds"` originalMap map[string]interface{} } @@ -246,7 +249,12 @@ func MessageFromMap(m interface{}) Message { } } - // TODO: Embeds + if iembeds, ok := mmap["embeds"]; ok { + embeds := iembeds.([]interface{}) + for _, iembed := range embeds { + msg.Embeds = append(msg.Embeds, EmbedFromMap(iembed)) + } + } return msg } @@ -277,15 +285,16 @@ func UserFromMap(m interface{}) User { return u } +// https://discord.com/developers/docs/resources/channel#attachment-object type Attachment struct { - ID string `json:"id"` - Filename string `json:"filename"` - ContentType string `json:"content_type"` - Size int `json:"size"` - Url string `json:"url"` - ProxyUrl string `json:"proxy_url"` - Height *int `json:"height"` - Width *int `json:"width"` + ID string `json:"id"` + Filename string `json:"filename"` + ContentType *string `json:"content_type"` + Size int `json:"size"` + Url string `json:"url"` + ProxyUrl string `json:"proxy_url"` + Height *int `json:"height"` + Width *int `json:"width"` } func AttachmentFromMap(m interface{}) Attachment { @@ -293,7 +302,7 @@ func AttachmentFromMap(m interface{}) Attachment { a := Attachment{ ID: mmap["id"].(string), Filename: mmap["filename"].(string), - ContentType: maybeString(mmap, "content_type"), + ContentType: maybeStringP(mmap, "content_type"), Size: int(mmap["size"].(float64)), Url: mmap["url"].(string), ProxyUrl: mmap["proxy_url"].(string), @@ -304,6 +313,224 @@ func AttachmentFromMap(m interface{}) Attachment { return a } +// https://discord.com/developers/docs/resources/channel#embed-object +type Embed struct { + Title *string `json:"title"` + Type *string `json:"type"` + Description *string `json:"description"` + Url *string `json:"url"` + Timestamp *string `json:"timestamp"` + Color *int `json:"color"` + Footer *EmbedFooter `json:"footer"` + Image *EmbedImage `json:"image"` + Thumbnail *EmbedThumbnail `json:"thumbnail"` + Video *EmbedVideo `json:"video"` + Provider *EmbedProvider `json:"provider"` + Author *EmbedAuthor `json:"author"` + Fields []EmbedField `json:"fields"` +} + +type EmbedFooter struct { + Text string `json:"text"` + IconUrl *string `json:"icon_url"` + ProxyIconUrl *string `json:"proxy_icon_url"` +} + +type EmbedImageish struct { + Url *string `json:"url"` + ProxyUrl *string `json:"proxy_url"` + Height *int `json:"height"` + Width *int `json:"width"` +} + +type EmbedImage struct { + EmbedImageish +} + +type EmbedThumbnail struct { + EmbedImageish +} + +type EmbedVideo struct { + EmbedImageish +} + +type EmbedProvider struct { + Name *string `json:"name"` + Url *string `json:"url"` +} + +type EmbedAuthor struct { + Name *string `json:"name"` + Url *string `json:"url"` + IconUrl *string `json:"icon_url"` + ProxyIconUrl *string `json:"proxy_icon_url"` +} + +type EmbedField struct { + Name string `json:"name"` + Value string `json:"value"` + Inline *bool `json:"inline"` +} + +func EmbedFromMap(m interface{}) Embed { + mmap := m.(map[string]interface{}) + + e := Embed{ + Title: maybeStringP(mmap, "title"), + Type: maybeStringP(mmap, "type"), + Description: maybeStringP(mmap, "description"), + Url: maybeStringP(mmap, "url"), + Timestamp: maybeStringP(mmap, "timestamp"), + Color: maybeIntP(mmap, "color"), + Footer: EmbedFooterFromMap(mmap, "footer"), + Image: EmbedImageFromMap(mmap, "image"), + Thumbnail: EmbedThumbnailFromMap(mmap, "thumbnail"), + Video: EmbedVideoFromMap(mmap, "video"), + Provider: EmbedProviderFromMap(mmap, "provider"), + Author: EmbedAuthorFromMap(mmap, "author"), + Fields: EmbedFieldsFromMap(mmap, "fields"), + } + + return e +} + +func EmbedFooterFromMap(m map[string]interface{}, k string) *EmbedFooter { + f, ok := m[k] + if !ok { + return nil + } + fMap, ok := f.(map[string]interface{}) + if !ok { + return nil + } + + return &EmbedFooter{ + Text: maybeString(fMap, "text"), + IconUrl: maybeStringP(fMap, "icon_url"), + ProxyIconUrl: maybeStringP(fMap, "proxy_icon_url"), + } +} + +func EmbedImageFromMap(m map[string]interface{}, k string) *EmbedImage { + val, ok := m[k] + if !ok { + return nil + } + valMap, ok := val.(map[string]interface{}) + if !ok { + return nil + } + + return &EmbedImage{ + EmbedImageish: EmbedImageish{ + Url: maybeStringP(valMap, "url"), + ProxyUrl: maybeStringP(valMap, "proxy_url"), + Height: maybeIntP(valMap, "height"), + Width: maybeIntP(valMap, "width"), + }, + } +} + +func EmbedThumbnailFromMap(m map[string]interface{}, k string) *EmbedThumbnail { + val, ok := m[k] + if !ok { + return nil + } + valMap, ok := val.(map[string]interface{}) + if !ok { + return nil + } + + return &EmbedThumbnail{ + EmbedImageish: EmbedImageish{ + Url: maybeStringP(valMap, "url"), + ProxyUrl: maybeStringP(valMap, "proxy_url"), + Height: maybeIntP(valMap, "height"), + Width: maybeIntP(valMap, "width"), + }, + } +} + +func EmbedVideoFromMap(m map[string]interface{}, k string) *EmbedVideo { + val, ok := m[k] + if !ok { + return nil + } + valMap, ok := val.(map[string]interface{}) + if !ok { + return nil + } + + return &EmbedVideo{ + EmbedImageish: EmbedImageish{ + Url: maybeStringP(valMap, "url"), + ProxyUrl: maybeStringP(valMap, "proxy_url"), + Height: maybeIntP(valMap, "height"), + Width: maybeIntP(valMap, "width"), + }, + } +} + +func EmbedProviderFromMap(m map[string]interface{}, k string) *EmbedProvider { + val, ok := m[k] + if !ok { + return nil + } + valMap, ok := val.(map[string]interface{}) + if !ok { + return nil + } + + return &EmbedProvider{ + Name: maybeStringP(valMap, "name"), + Url: maybeStringP(valMap, "url"), + } +} + +func EmbedAuthorFromMap(m map[string]interface{}, k string) *EmbedAuthor { + val, ok := m[k] + if !ok { + return nil + } + valMap, ok := val.(map[string]interface{}) + if !ok { + return nil + } + + return &EmbedAuthor{ + Name: maybeStringP(valMap, "name"), + Url: maybeStringP(valMap, "url"), + } +} + +func EmbedFieldsFromMap(m map[string]interface{}, k string) []EmbedField { + val, ok := m[k] + if !ok { + return nil + } + valSlice, ok := val.([]interface{}) + if !ok { + return nil + } + + var result []EmbedField + for _, innerVal := range valSlice { + valMap, ok := innerVal.(map[string]interface{}) + if !ok { + continue + } + + result = append(result, EmbedField{ + Name: maybeString(valMap, "name"), + Value: maybeString(valMap, "value"), + Inline: maybeBoolP(valMap, "inline"), + }) + } + + return result +} + func maybeString(m map[string]interface{}, k string) string { val, ok := m[k] if !ok { @@ -337,3 +564,20 @@ func maybeIntP(m map[string]interface{}, k string) *int { intval := int(val.(float64)) return &intval } + +func maybeBool(m map[string]interface{}, k string) bool { + val, ok := m[k] + if !ok { + return false + } + return val.(bool) +} + +func maybeBoolP(m map[string]interface{}, k string) *bool { + val, ok := m[k] + if !ok { + return nil + } + boolval := val.(bool) + return &boolval +} diff --git a/src/discord/showcase.go b/src/discord/showcase.go index 65abbdc..a0de713 100644 --- a/src/discord/showcase.go +++ b/src/discord/showcase.go @@ -3,17 +3,20 @@ package discord import ( "context" "errors" + "fmt" "io" "net/http" "net/url" "regexp" "strings" + "time" "git.handmade.network/hmn/hmn/src/assets" "git.handmade.network/hmn/hmn/src/db" "git.handmade.network/hmn/hmn/src/logging" "git.handmade.network/hmn/hmn/src/models" "git.handmade.network/hmn/hmn/src/oops" + "github.com/google/uuid" "github.com/jackc/pgx/v4" ) @@ -21,7 +24,8 @@ var reDiscordMessageLink = regexp.MustCompile(`https?://.+?(\s|$)`) var errNotEnoughInfo = errors.New("Discord didn't send enough info in this event for us to do this") -func (bot *discordBotInstance) processShowcaseMsg(ctx context.Context, msg *Message) error { +// TODO: Can this function be called asynchronously? +func (bot *botInstance) processShowcaseMsg(ctx context.Context, msg *Message) error { switch msg.Type { case MessageTypeDefault, MessageTypeReply, MessageTypeApplicationCommand: default: @@ -69,7 +73,7 @@ func (bot *discordBotInstance) processShowcaseMsg(ctx context.Context, msg *Mess return nil } -func (bot *discordBotInstance) maybeDeleteShowcaseMsg(ctx context.Context, msg *Message) (didDelete bool, err error) { +func (bot *botInstance) maybeDeleteShowcaseMsg(ctx context.Context, msg *Message) (didDelete bool, err error) { hasGoodContent := true if msg.OriginalHasFields("content") && !messageHasLinks(msg.Content) { hasGoodContent = false @@ -116,7 +120,7 @@ the database. This does not create snippets or do anything besides save the message itself. */ -func (bot *discordBotInstance) saveMessage( +func (bot *botInstance) saveMessage( ctx context.Context, tx pgx.Tx, msg *Message, @@ -180,7 +184,7 @@ snippets. Idempotent; can be called any time whether the message exists or not. */ -func (bot *discordBotInstance) saveMessageAndContents( +func (bot *botInstance) saveMessageAndContents( ctx context.Context, tx pgx.Tx, msg *Message, @@ -230,12 +234,53 @@ func (bot *discordBotInstance) saveMessageAndContents( } } - // TODO: Save embeds + // Save embeds + for _, embed := range msg.Embeds { + _, err := bot.saveEmbed(ctx, tx, &embed, discordUser.HMNUserId, msg.ID) + if err != nil { + return nil, oops.New(err, "failed to save embed") + } + } return newMsg, nil } -func (bot *discordBotInstance) saveAttachment(ctx context.Context, tx pgx.Tx, attachment *Attachment, hmnUserID int, discordMessageID string) (*models.DiscordMessageAttachment, error) { +var discordDownloadClient = &http.Client{ + Timeout: 10 * time.Second, +} + +type DiscordResourceBadStatusCode error + +func downloadDiscordResource(ctx context.Context, url string) ([]byte, string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil) + if err != nil { + return nil, "", oops.New(err, "failed to make Discord download request") + } + res, err := discordDownloadClient.Do(req) + if err != nil { + return nil, "", oops.New(err, "failed to fetch Discord resource data") + } + defer res.Body.Close() + + if res.StatusCode < 200 || 299 < res.StatusCode { + return nil, "", DiscordResourceBadStatusCode(fmt.Errorf("status code %d from Discord resource: %s", res.StatusCode, url)) + } + + content, err := io.ReadAll(res.Body) + if err != nil { + panic(err) + } + + return content, res.Header.Get("Content-Type"), nil +} + +func (bot *botInstance) saveAttachment( + ctx context.Context, + tx pgx.Tx, + attachment *Attachment, + hmnUserID int, + discordMessageID string, +) (*models.DiscordMessageAttachment, error) { // TODO: Return an existing attachment if it exists width := 0 @@ -247,22 +292,20 @@ func (bot *discordBotInstance) saveAttachment(ctx context.Context, tx pgx.Tx, at height = *attachment.Height } - // TODO: Timeouts and stuff, context cancellation - res, err := http.Get(attachment.Url) + content, _, err := downloadDiscordResource(ctx, attachment.Url) if err != nil { - return nil, oops.New(err, "failed to fetch attachment data") + return nil, oops.New(err, "failed to download Discord attachment") } - defer res.Body.Close() - content, err := io.ReadAll(res.Body) - if err != nil { - panic(err) + contentType := "application/octet-stream" + if attachment.ContentType != nil { + contentType = *attachment.ContentType } asset, err := assets.Create(ctx, tx, assets.CreateInput{ - Content: content, - Filename: attachment.Filename, - MimeType: attachment.ContentType, + Content: content, + Filename: attachment.Filename, + ContentType: contentType, UploaderID: &hmnUserID, Width: width, @@ -301,7 +344,109 @@ func (bot *discordBotInstance) saveAttachment(ctx context.Context, tx pgx.Tx, at return iDiscordAttachment.(*models.DiscordMessageAttachment), nil } -func (bot *discordBotInstance) allowedToCreateMessageSnippet(ctx context.Context, msg *Message) (bool, error) { +func (bot *botInstance) saveEmbed( + ctx context.Context, + tx pgx.Tx, + embed *Embed, + hmnUserID int, + discordMessageID string, +) (*models.DiscordMessageEmbed, error) { + // TODO: Does this need to be idempotent + + isOkImageType := func(contentType string) bool { + return strings.HasPrefix(contentType, "image/") + } + + isOkVideoType := func(contentType string) bool { + return strings.HasPrefix(contentType, "video/") + } + + maybeSaveImageish := func(i EmbedImageish, contentTypeCheck func(string) bool) (*uuid.UUID, error) { + content, contentType, err := downloadDiscordResource(ctx, *i.Url) + if err != nil { + var statusError DiscordResourceBadStatusCode + if errors.As(err, &statusError) { + return nil, nil + } else { + return nil, oops.New(err, "failed to save Discord embed") + } + } + if contentTypeCheck(contentType) { + in := assets.CreateInput{ + Content: content, + Filename: "embed", + ContentType: contentType, + UploaderID: &hmnUserID, + } + + if i.Width != nil { + in.Width = *i.Width + } + if i.Height != nil { + in.Height = *i.Height + } + + asset, err := assets.Create(ctx, tx, in) + if err != nil { + return nil, oops.New(err, "failed to create asset from embed") + } + return &asset.ID, nil + } + + return nil, nil + } + + var imageAssetId *uuid.UUID + var videoAssetId *uuid.UUID + var err error + + if embed.Video != nil && embed.Video.Url != nil { + videoAssetId, err = maybeSaveImageish(embed.Video.EmbedImageish, isOkVideoType) + } else if embed.Image != nil && embed.Image.Url != nil { + imageAssetId, err = maybeSaveImageish(embed.Image.EmbedImageish, isOkImageType) + } else if embed.Thumbnail != nil && embed.Thumbnail.Url != nil { + imageAssetId, err = maybeSaveImageish(embed.Thumbnail.EmbedImageish, isOkImageType) + } + if err != nil { + return nil, err + } + + // Save the embed into the db + // TODO(db): Insert, RETURNING + var savedEmbedId int + err = tx.QueryRow(ctx, + ` + INSERT INTO handmade_discordmessageembed (title, description, url, message_id, image_id, video_id) + VALUES ($1, $2, $3, $4, $5, $6) + RETURNING id + `, + embed.Title, + embed.Description, + embed.Url, + discordMessageID, + imageAssetId, + videoAssetId, + ).Scan(&savedEmbedId) + if err != nil { + return nil, oops.New(err, "failed to insert new embed") + } + + iDiscordEmbed, err := db.QueryOne(ctx, tx, models.DiscordMessageEmbed{}, + ` + SELECT $columns + FROM handmade_discordmessageembed + WHERE id = $1 + `, + savedEmbedId, + ) + if err != nil { + return nil, oops.New(err, "failed to fetch new Discord embed data") + } + + return iDiscordEmbed.(*models.DiscordMessageEmbed), nil +} + +func (bot *botInstance) allowedToCreateMessageSnippet(ctx context.Context, msg *Message) (bool, error) { canSave, err := db.QueryBool(ctx, bot.dbConn, ` SELECT u.discord_save_showcase @@ -322,7 +467,7 @@ func (bot *discordBotInstance) allowedToCreateMessageSnippet(ctx context.Context return canSave, nil } -func (bot *discordBotInstance) createMessageSnippet(ctx context.Context, msg *Message) (*models.Snippet, error) { +func (bot *botInstance) createMessageSnippet(ctx context.Context, msg *Message) (*models.Snippet, error) { // TODO: Actually do this return nil, nil }