Save Discord attachments and embeds

This commit is contained in:
Ben Visness 2021-08-23 16:52:57 -05:00
parent 76f9256e97
commit 72ae938302
5 changed files with 436 additions and 47 deletions

View File

@ -50,7 +50,7 @@ func init() {
type CreateInput struct {
Content []byte
Filename string
MimeType 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,

View File

@ -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

View File

@ -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:

View File

@ -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,10 +285,11 @@ 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"`
ContentType *string `json:"content_type"`
Size int `json:"size"`
Url string `json:"url"`
ProxyUrl string `json:"proxy_url"`
@ -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
}

View File

@ -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,
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
}