diff --git a/Sources/Encryption/SenderKeyRecord.swift b/Sources/Encryption/SenderKeyRecord.swift
index 175b64a..4813c44 100644
--- a/Sources/Encryption/SenderKeyRecord.swift
+++ b/Sources/Encryption/SenderKeyRecord.swift
@@ -44,7 +44,7 @@ final class SenderKeyRecord {
     func state(for id: UInt32) -> SenderKeyState? {
         for item in states {
             if item.keyId == id {
-                return state
+                return item
             }
         }
         return nil
diff --git a/Tests/SenderKeyRecordTests.swift b/Tests/SenderKeyRecordTests.swift
index ac09671..865f380 100644
--- a/Tests/SenderKeyRecordTests.swift
+++ b/Tests/SenderKeyRecordTests.swift
@@ -109,6 +109,18 @@ class SenderKeyRecordTests: XCTestCase {
             chainKey: senderKey2,
             signatureKeyPair: senderSigningKey2)
 
+        guard let state = record.state(for: stateId1) else {
+            XCTFail("\(stateId1) missed")
+            return
+        }
+        XCTAssertEqual(state.keyId, stateId1)
+
+        guard let state = record.state(for: stateId2) else {
+            XCTFail("\(stateId2) missed")
+            return
+        }
+        XCTAssertEqual(state.keyId, stateId2)
+        
         guard let serialized = try? record.protoData() else {
             XCTFail("Could not serialize SenderKeyRecord")
             return