From 008d025fa02ce895b0e0f147c61e60fd02b051cd Mon Sep 17 00:00:00 2001
From: William Perron <hey@wperron.io>
Date: Mon, 1 Jan 2024 16:05:07 -0500
Subject: [PATCH] apply correlation IDs to modal submits and message actions

---
 cmd/themis-server/main.go | 84 +++++++++++++++++++++++++++++++++------
 1 file changed, 71 insertions(+), 13 deletions(-)

diff --git a/cmd/themis-server/main.go b/cmd/themis-server/main.go
index 477629f..fa505e1 100644
--- a/cmd/themis-server/main.go
+++ b/cmd/themis-server/main.go
@@ -7,6 +7,7 @@ import (
 	"flag"
 	"fmt"
 	"net/http"
+	"net/url"
 	"os"
 	"os/signal"
 	"sort"
@@ -465,10 +466,22 @@ func main() {
 			}
 		},
 		"flush": func(ctx context.Context, s *discordgo.Session, i *discordgo.InteractionCreate) {
+			cid := correlation.FromContext(ctx)
+			baggage := make(url.Values)
+			baggage.Set("correlation_id", cid.String())
+			state := baggage.Encode()
+
+			sb := strings.Builder{}
+			sb.WriteString("modal_flush")
+			if state != "" {
+				sb.WriteRune(':')
+				sb.WriteString(state)
+			}
+
 			if err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
 				Type: discordgo.InteractionResponseModal,
 				Data: &discordgo.InteractionResponseData{
-					CustomID: "modals_flush_" + i.Interaction.Member.User.ID,
+					CustomID: sb.String(),
 					Title:    "Are you sure?",
 					Components: []discordgo.MessageComponent{
 						discordgo.ActionsRow{
@@ -740,28 +753,42 @@ func registerHandlers(sess *discordgo.Session, handlers map[string]Handler) {
 		log.Info().Str("user_id", fmt.Sprintf("%s#%s", s.State.User.Username, s.State.User.Discriminator)).Msg("logged in")
 	})
 	sess.AddHandler(func(s *discordgo.Session, i *discordgo.InteractionCreate) {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-		ctx = context.WithValue(ctx, "correlation_id", gen.Next())
-
 		switch i.Type {
 		case discordgo.InteractionApplicationCommand:
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+			ctx = context.WithValue(ctx, "correlation_id", gen.Next())
+
 			if h, ok := handlers[i.ApplicationCommandData().Name]; ok {
 				withLogging(i.ApplicationCommandData().Name, h)(ctx, s, i)
 			}
-		// TODO(wperron) apply correlation IDs to the other interation types
-		// TODO(wperron) is it possible to correlate a modal submit or message
-		//               component to the original interaction?
 		case discordgo.InteractionModalSubmit:
-			if strings.HasPrefix(i.ModalSubmitData().CustomID, "modals_flush_") {
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+
+			state, err := parseCustomIDState(i.ModalSubmitData().CustomID)
+			if err != nil {
+				log.Error().Ctx(ctx).Err(err).Msg("unexpected error occured while parsing state from custom id, returning early.")
+				return
+			}
+
+			cid := state.Get("correlation_id")
+			if cid != "" {
+				ctx = context.WithValue(ctx, "correlation_id", cid)
+			} else {
+				ctx = context.WithValue(ctx, "correlation_id", gen.Next())
+			}
+
+			if strings.HasPrefix(i.ModalSubmitData().CustomID, "modal_flush") {
 				sub := i.ModalSubmitData().Components[0].(*discordgo.ActionsRow).Components[0].(*discordgo.TextInput).Value
 				sub = strings.ToLower(sub)
-				log.Debug().Str("value", sub).Msg("flush modal submitted")
+				log.Debug().Ctx(ctx).Str("value", sub).Msg("flush modal submitted")
+
 				if sub == "y" || sub == "ye" || sub == "yes" {
 					err := store.Flush(context.Background(), i.Member.User.ID)
 					msg := "Flushed all claims!"
 					if err != nil {
-						log.Error().Err(err).Msg("failed to flush claims")
+						log.Error().Ctx(ctx).Err(err).Msg("failed to flush claims")
 						msg = "failed to flush claims from database"
 					}
 
@@ -772,7 +799,7 @@ func registerHandlers(sess *discordgo.Session, handlers map[string]Handler) {
 						},
 					})
 					if err != nil {
-						log.Error().Err(err).Msg("failed to respond to interaction")
+						log.Error().Ctx(ctx).Err(err).Msg("failed to respond to interaction")
 					}
 					return
 				}
@@ -784,14 +811,32 @@ func registerHandlers(sess *discordgo.Session, handlers map[string]Handler) {
 					},
 				})
 				if err != nil {
-					log.Error().Err(err).Msg("failed to respond to interaction")
+					log.Error().Ctx(ctx).Err(err).Msg("failed to respond to interaction")
 				}
 				return
 			}
 		case discordgo.InteractionMessageComponent:
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+
+			state, err := parseCustomIDState(i.MessageComponentData().CustomID)
+			if err != nil {
+				log.Error().Ctx(ctx).Err(err).Msg("unexpected error occured while parsing state from custom id, returning early.")
+				return
+			}
+
+			cid := state.Get("correlation_id")
+			if cid != "" {
+				ctx = context.WithValue(ctx, "correlation_id", cid)
+			} else {
+				ctx = context.WithValue(ctx, "correlation_id", gen.Next())
+			}
+
 			switch i.MessageComponentData().CustomID {
 			case "schedule-response":
 				userId := i.Member.User.ID
+				log.Info().Ctx(ctx).Str("message_component", "schedule-response").Str("userid", userId).Msg("handling message component interaction")
+
 				if err := store.AddAbsence(context.TODO(), themis.NextMonday(), userId); err != nil {
 					if err := s.InteractionRespond(i.Interaction, &discordgo.InteractionResponse{
 						Type: discordgo.InteractionResponseChannelMessageWithSource,
@@ -949,3 +994,16 @@ func min(a, b int) int {
 	}
 	return b
 }
+
+func parseCustomIDState(qs string) (url.Values, error) {
+	parts := strings.Split(qs, ":")
+	if len(parts) == 1 {
+		return make(url.Values), nil
+	}
+
+	v, err := url.ParseQuery(parts[1])
+	if err != nil {
+		return nil, err
+	}
+	return v, nil
+}