fix: keep E2E ciphertext and is_encrypted correct on message edits

Author: Cod-e-Codes

ARCHITECTURE.md

@@ -274,6 +274,7 @@ See [PROTOCOL.md](PROTOCOL.md) for the full message format specification.
 3. **Message encryption**: Inner JSON payload is sealed with ChaCha20-Poly1305; nonce ‖ ciphertext is base64-encoded into `content` with `encrypted: true` (see **PROTOCOL.md**)
 4. **Transport**: Standard WebSocket JSON; server does not decrypt
 5. **Storage**: Server persists opaque ciphertext in the database
+6. **Edits**: For `type: edit`, the reference client sends the replacement body ciphertext when E2E is enabled; the server updates `content` and persists `is_encrypted` from the wire `encrypted` field so reconnects and message metadata stay correct
 
 ## Database Schema
 
@@ -299,6 +300,8 @@ CREATE TABLE messages (
 );
 ```
 
+Edits update `content` and `edited`; `is_encrypted` follows the `encrypted` field on the edit message so ciphertext is not downgraded to plain text in the database after an edit.
+
 #### `user_message_state`
 ```sql
 CREATE TABLE user_message_state (

PROTOCOL.md

@@ -122,7 +122,7 @@ These values of `type` extend the core chat protocol:
 
 | `type` | Purpose |
 |--------|---------|
-| `edit` | Replace the text of an existing message. Requires `message_id` and new text in `content`. Only the original author may edit (server-enforced). |
+| `edit` | Replace the text of an existing message. Requires `message_id` and new body in `content`. Set `encrypted` to `true` when `content` is E2E ciphertext (same base64 **nonce ‖ ciphertext** layout as `text`); the server persists that flag to `is_encrypted` so history and reconnects stay consistent. Only the original sender may edit; admins cannot edit someone else's message (unlike `delete`, where an admin may remove any message). Enforced by matching WebSocket `sender` to the stored row. |
 | `delete` | Soft-delete a message. Requires `message_id`. Authors may delete their own messages; admins may delete any message. |
 | `typing` | Typing indicator; `content` is not required. The server sets `sender` and broadcasts to clients (see [Channels](#channels) for delivery scope when `channel` is set). |
 | `reaction` | Add or remove a reaction. Requires the `reaction` object (`emoji`, `target_id`, optional `is_removal`). |
@@ -158,6 +158,7 @@ Optional chat encryption uses a **shared symmetric key** for all participants (g
 - **Algorithm**: ChaCha20-Poly1305 (RFC 8439). Each encrypted payload uses a random **12-byte nonce** (typical for this AEAD).
 - **Text messages on the wire**: Same JSON message shape as plaintext chat; set `encrypted` to `true`. The `content` field is **standard base64** encoding of **nonce ‖ ciphertext** (nonce first, 12 bytes, then the Poly1305-sealed ciphertext). The plaintext decrypted by the AEAD is a JSON object representing the inner chat message (e.g. sender, content, type, timestamp) as produced by the reference client.
 - **Files**: When `type` is `file` and E2E is enabled, the reference client encrypts file bytes with the same global key; see client implementation for the exact binary layout (nonce-prefixed ciphertext).
+- **Edits** (`type`: `edit`): When E2E is enabled, the reference client encrypts the new plaintext the same way as a normal `text` message and sets `encrypted` accordingly. The server stores the new opaque `content` and updates `is_encrypted` from the incoming `encrypted` field (it does not force plaintext).
 
 The server stores and relays opaque `content` (and encrypted file blobs) without performing decryption.
 
@@ -175,6 +176,7 @@ The server stores and relays opaque `content` (and encrypted file blobs) without
   - Delivers to all connected clients **or** only to members of a channel when `channel` is non-empty and `sender` is not `System` (see [Channels](#channels)). Direct messages use a separate path (sender and recipient only).
 - Reactions, read receipts, and last channel per user may be persisted server-side and replayed to reconnecting clients.
 - Message history is capped at 1000 messages.
+- On successful `type`: `edit`, the server updates `content`, sets `edited`, and sets stored encryption metadata from the incoming `encrypted` field (so edited ciphertext rows remain ciphertext with `is_encrypted` aligned to the wire).
 
 ---
 

README.md

@@ -347,7 +347,7 @@ Run **`./marchat-client -doctor`** or **`./marchat-server -doctor`** for a text
 ### Messaging
 | Command | Description |
 |---------|-------------|
-| `:edit <id> <text>` | Edit a message by its ID |
+| `:edit <id> <text>` | Edit your own message by ID (admins cannot edit others' messages; with E2E on, the new text is encrypted like normal chat and the server keeps `is_encrypted` in sync) |
 | `:delete <id>` | Delete a message by its ID |
 | `:dm [user] [msg]` | Send a DM or toggle DM mode (no args exits DM mode) |
 | `:search <query>` | Search message history on the server |

client/main.go

@@ -595,6 +595,7 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				if m2.MessageID == v.MessageID {
 					m.messages[i].Content = v.Content
 					m.messages[i].Edited = true
+					m.messages[i].Encrypted = v.Encrypted
 					break
 				}
 			}
@@ -1560,8 +1561,23 @@ func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 						m.messages = append(m.messages, shared.Message{Sender: "System", Content: "Invalid message ID", CreatedAt: time.Now(), Type: shared.TextMessage})
 					} else {
 						newContent := strings.Join(parts[2:], " ")
-						editMsg := shared.Message{Type: shared.EditMessageType, MessageID: id, Content: newContent, Sender: m.cfg.Username}
-						if m.conn != nil {
+						editMsg := shared.Message{Type: shared.EditMessageType, MessageID: id, Sender: m.cfg.Username}
+						okToSend := true
+						if m.useE2E && m.keystore != nil && m.keystore.GetSessionKey("global") != nil {
+							if err := verifyKeystoreUnlocked(m.keystore); err != nil {
+								m.messages = append(m.messages, shared.Message{Sender: "System", Content: "[ERROR] Keystore not unlocked: " + err.Error(), CreatedAt: time.Now(), Type: shared.TextMessage})
+								okToSend = false
+							} else if wire, encErr := encryptGlobalTextWireContent(m.keystore, m.cfg.Username, newContent); encErr != nil {
+								m.messages = append(m.messages, shared.Message{Sender: "System", Content: "[ERROR] Failed to encrypt edit: " + encErr.Error(), CreatedAt: time.Now(), Type: shared.TextMessage})
+								okToSend = false
+							} else {
+								editMsg.Content = wire
+								editMsg.Encrypted = true
+							}
+						} else {
+							editMsg.Content = newContent
+						}
+						if okToSend && m.conn != nil {
 							_ = m.conn.WriteJSON(editMsg)
 						}
 					}

client/websocket.go

@@ -39,6 +39,32 @@ type UserList struct {
 
 type quitMsg struct{}
 
+// encryptGlobalTextWireContent returns base64(nonce ‖ ciphertext) for global chat E2E text,
+// matching the wire format produced for normal encrypted messages.
+func encryptGlobalTextWireContent(keystore *crypto.KeyStore, username, plaintext string) (string, error) {
+	if keystore == nil {
+		return "", fmt.Errorf("keystore not initialized")
+	}
+	if keystore.GetSessionKey("global") == nil {
+		return "", fmt.Errorf("global key not available - global E2E encryption not initialized")
+	}
+	encryptedMsg, err := keystore.EncryptMessage(username, plaintext, "global")
+	if err != nil {
+		return "", fmt.Errorf("global encryption failed: %w", err)
+	}
+	if len(encryptedMsg.Encrypted) == 0 {
+		return "", fmt.Errorf("encryption returned empty ciphertext")
+	}
+	combinedData := make([]byte, 0, len(encryptedMsg.Nonce)+len(encryptedMsg.Encrypted))
+	combinedData = append(combinedData, encryptedMsg.Nonce...)
+	combinedData = append(combinedData, encryptedMsg.Encrypted...)
+	finalContent := base64.StdEncoding.EncodeToString(combinedData)
+	if len(finalContent) == 0 {
+		return "", fmt.Errorf("final content is empty after encoding")
+	}
+	return finalContent, nil
+}
+
 func debugEncryptAndSend(recipients []string, plaintext string, ws *websocket.Conn, keystore *crypto.KeyStore, username string) error {
 	log.Printf("DEBUG: Starting global encryption for %d recipients", len(recipients))
 	log.Printf("DEBUG: Plaintext length: %d", len(plaintext))
@@ -56,32 +82,13 @@ func debugEncryptAndSend(recipients []string, plaintext string, ws *websocket.Co
 	}
 	log.Printf("DEBUG: Global key available (ID: %s)", globalKey.KeyID)
 
-	conversationID := "global"
-	encryptedMsg, err := keystore.EncryptMessage(username, plaintext, conversationID)
+	finalContent, err := encryptGlobalTextWireContent(keystore, username, plaintext)
 	if err != nil {
 		log.Printf("ERROR: Global encryption failed: %v", err)
-		return fmt.Errorf("global encryption failed: %v", err)
-	}
-
-	log.Printf("DEBUG: Global encryption successful - encrypted length: %d", len(encryptedMsg.Encrypted))
-
-	if len(encryptedMsg.Encrypted) == 0 {
-		log.Printf("ERROR: Encryption returned empty ciphertext")
-		return fmt.Errorf("encryption returned empty ciphertext; aborting send")
+		return err
 	}
-
-	combinedData := make([]byte, 0, len(encryptedMsg.Nonce)+len(encryptedMsg.Encrypted))
-	combinedData = append(combinedData, encryptedMsg.Nonce...)
-	combinedData = append(combinedData, encryptedMsg.Encrypted...)
-
-	finalContent := base64.StdEncoding.EncodeToString(combinedData)
 	log.Printf("DEBUG: Base64 encoded nonce+ciphertext - length: %d", len(finalContent))
 
-	if len(finalContent) == 0 {
-		log.Printf("ERROR: Final content is empty after encoding")
-		return fmt.Errorf("final content is empty after encoding")
-	}
-
 	msg := shared.Message{
 		Content:   finalContent,
 		Sender:    username,

server/client.go

@@ -107,7 +107,7 @@ func (c *Client) readPump() {
 		}
 
 		if msg.Type == shared.EditMessageType && msg.MessageID > 0 {
-			if err := EditMessage(c.db, msg.MessageID, c.username, msg.Content); err != nil {
+			if err := EditMessage(c.db, msg.MessageID, c.username, msg.Content, msg.Encrypted); err != nil {
 				c.send <- shared.Message{
 					Sender:    "System",
 					Content:   "Edit failed: " + err.Error(),

server/handlers.go

@@ -697,10 +697,10 @@ func ClearMessages(db *sql.DB) error {
 	return err
 }
 
-func EditMessage(db *sql.DB, messageID int64, sender, newContent string) error {
+func EditMessage(db *sql.DB, messageID int64, sender, newContent string, encrypted bool) error {
 	result, err := dbExec(db,
-		`UPDATE messages SET content = ?, edited = 1, is_encrypted = 0 WHERE message_id = ? AND sender = ?`,
-		newContent, messageID, sender)
+		`UPDATE messages SET content = ?, edited = 1, is_encrypted = ? WHERE message_id = ? AND sender = ?`,
+		newContent, encrypted, messageID, sender)
 	if err != nil {
 		return fmt.Errorf("edit message: %w", err)
 	}

server/handlers_test.go

@@ -458,12 +458,13 @@ func TestEditMessage(t *testing.T) {
 	}
 
 	t.Run("success", func(t *testing.T) {
-		if err := EditMessage(db, id, "alice", "updated"); err != nil {
+		if err := EditMessage(db, id, "alice", "updated", false); err != nil {
 			t.Fatalf("EditMessage failed: %v", err)
 		}
 		var content string
 		var edited bool
-		if err := db.QueryRow(`SELECT content, edited FROM messages WHERE message_id = ?`, id).Scan(&content, &edited); err != nil {
+		var isEnc bool
+		if err := db.QueryRow(`SELECT content, edited, is_encrypted FROM messages WHERE message_id = ?`, id).Scan(&content, &edited, &isEnc); err != nil {
 			t.Fatalf("query row: %v", err)
 		}
 		if content != "updated" {
@@ -472,6 +473,35 @@ func TestEditMessage(t *testing.T) {
 		if !edited {
 			t.Error("edited flag not set")
 		}
+		if isEnc {
+			t.Error("plain message edit should leave is_encrypted false")
+		}
+	})
+
+	t.Run("encrypted_content_flag", func(t *testing.T) {
+		eid, err := InsertMessage(db, shared.Message{
+			Sender:    "alice",
+			Content:   "original-cipher",
+			CreatedAt: now.Add(2 * time.Minute),
+			Encrypted: true,
+		})
+		if err != nil {
+			t.Fatalf("InsertMessage failed: %v", err)
+		}
+		if err := EditMessage(db, eid, "alice", "updated-cipher", true); err != nil {
+			t.Fatalf("EditMessage failed: %v", err)
+		}
+		var content string
+		var isEnc bool
+		if err := db.QueryRow(`SELECT content, is_encrypted FROM messages WHERE message_id = ?`, eid).Scan(&content, &isEnc); err != nil {
+			t.Fatalf("query row: %v", err)
+		}
+		if content != "updated-cipher" {
+			t.Errorf("content = %q, want updated-cipher", content)
+		}
+		if !isEnc {
+			t.Error("encrypted edit should set is_encrypted true in DB")
+		}
 	})
 
 	t.Run("wrong user", func(t *testing.T) {
@@ -484,7 +514,7 @@ func TestEditMessage(t *testing.T) {
 		if err != nil {
 			t.Fatalf("InsertMessage failed: %v", err)
 		}
-		err = EditMessage(db, id2, "alice", "hijack")
+		err = EditMessage(db, id2, "alice", "hijack", false)
 		if err == nil {
 			t.Fatal("expected error for wrong sender")
 		}