It works! 🎉

supercaracal · Jun 26, 2021 · 8a48658 · 8a48658
1 parent 9b5a676
commit 8a48658
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 17 deletions.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ $ make
 ```
 $ ./encrypt
 Raw password:
-SCRAM-SHA-256$4096:yTo5lMI+1XyqZOcvYz99Kw==$VJcML25bB3h0xiMUFw9D4spAJwp8IxD1CxnkR7XPty8=:NE05auswTZk1ntaXa8DrO9tYekyhfv1qRMXmugXpGPc=
+SCRAM-SHA-256$4096:Mg8UNqSaPstxvBVRVYPQTw==$Zl7Rhln+rus3z+4YwC+7CgL/uKSUvqWH8mHMUizh1EI=:G9dSawW20CNLxTnZdcwHEHg9U9hG2noNEV2/t7ptq3s=
 ```
 
 ### Testing
@@ -19,7 +19,9 @@ $ docker run --rm --name=test -e POSTGRES_PASSWORD=postgres -e POSTGRES_INITDB_A
 ```
 
 ```
-$ docker exec -it test psql -U postgres -c "CREATE ROLE test WITH LOGIN PASSWORD 'SCRAM-SHA-256$4096:yTo5lMI+1XyqZOcvYz99Kw==$VJcML25bB3h0xiMUFw9D4spAJwp8IxD1CxnkR7XPty8=:NE05auswTZk1ntaXa8DrO9tYekyhfv1qRMXmugXpGPc='"
+$ docker exec -it test bash -c 'cat | psql -U postgres'
+CREATE ROLE test WITH LOGIN PASSWORD 'SCRAM-SHA-256$4096:Mg8UNqSaPstxvBVRVYPQTw==$Zl7Rhln+rus3z+4YwC+7CgL/uKSUvqWH8mHMUizh1EI=:G9dSawW20CNLxTnZdcwHEHg9U9hG2noNEV2/t7ptq3s='
+## press Ctrl-D
 CREATE ROLE
 ```
 
@@ -28,7 +30,7 @@ $ docker exec -it test psql -U postgres -c 'SELECT usename, passwd FROM pg_catal
  usename  |                                                                passwd                                                                 
 ----------+---------------------------------------------------------------------------------------------------------------------------------------
  postgres | SCRAM-SHA-256$4096:N+t+PZUQAu25roNaMJiQIw==$MNmcJjqjLwfWBTvKq2zRCWSWPFQX6KnDqqyrqA1XU5g=:jL3qX7jzS4wSP1rOmEbbmLReYL98WeKukK8SfLcdpvU= 
- test     | SCRAM-SHA-256$4096:yTo5lMI+1XyqZOcvYz99Kw==$VJcML25bB3h0xiMUFw9D4spAJwp8IxD1CxnkR7XPty8=:NE05auswTZk1ntaXa8DrO9tYekyhfv1qRMXmugXpGPc= 
+ test     | SCRAM-SHA-256$4096:Mg8UNqSaPstxvBVRVYPQTw==$Zl7Rhln+rus3z+4YwC+7CgL/uKSUvqWH8mHMUizh1EI=:G9dSawW20CNLxTnZdcwHEHg9U9hG2noNEV2/t7ptq3s= 
 (2 rows)
 ```
 

diff --git a/main.go b/main.go
@@ -21,8 +21,13 @@ import (
 )
 
 const (
-	saltSize     = 16
-	digestLen    = 32
+	// @see https://github.com/postgres/postgres/blob/e6bdfd9700ebfc7df811c97c2fc46d7e94e329a2/src/include/common/scram-common.h#L36-L41
+	saltSize = 16
+
+	// @see https://github.com/postgres/postgres/blob/c30f54ad732ca5c8762bb68bbe0f51de9137dd72/src/include/common/sha2.h#L22
+	digestLen = 32
+
+	// @see https://github.com/postgres/postgres/blob/e6bdfd9700ebfc7df811c97c2fc46d7e94e329a2/src/include/common/scram-common.h#L43-L47
 	iterationCnt = 4096
 )
 
@@ -48,27 +53,23 @@ func encodeB64(src []byte) (dst []byte) {
 	return
 }
 
-func getHashSum(key, msg []byte) []byte {
+func getHMACSum(key, msg []byte) []byte {
 	h := hmac.New(sha256.New, key)
 	_, _ = h.Write(msg)
 	return h.Sum(nil)
 }
 
-func getStoredKey(clientKey []byte) (storedKey []byte) {
-	key := sha256.Sum256(clientKey)
-	storedKey = make([]byte, 0, len(key))
-	for _, b := range key {
-		storedKey = append(storedKey, b)
-	}
-	return
+func getSHA256Sum(key []byte) []byte {
+	h := sha256.New()
+	_, _ = h.Write(key)
+	return h.Sum(nil)
 }
 
-// FIXME: some bugs
 func encryptPassword(rawPassword, salt []byte, iter, keyLen int) string {
 	digestKey := pbkdf2.Key(rawPassword, salt, iter, keyLen, sha256.New)
-	clientKey := getHashSum(digestKey, []byte("Client Key"))
-	storedKey := getStoredKey(clientKey)
-	serverKey := getHashSum(digestKey, []byte("Server Key"))
+	clientKey := getHMACSum(digestKey, []byte("Client Key"))
+	storedKey := getSHA256Sum(clientKey)
+	serverKey := getHMACSum(digestKey, []byte("Server Key"))
 
 	return fmt.Sprintf("SCRAM-SHA-256$%d:%s$%s:%s",
 		iter,