Refactored detection logic to include ports and a confidence score

protocol/detect_postgres.go

@@ -20,16 +20,20 @@ func registerPostgreSQL() {
 	Register(Client, "????\x00\x03\x00\x00", detectPostgreSQLClient) // Startup packet, protocol 3.0
 }
 
-func detectPostgreSQLClient(dir Direction, data []byte) *Protocol {
+func detectPostgreSQLClient(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64) {
 	// A client startup message needs at least 8 bytes (length + protocol version).
 	if len(data) < 8 {
-		return nil
+		return nil, 0
 	}
 
 	length := int(binary.BigEndian.Uint32(data[0:]))
 	if len(data) != length {
 		log.Printf("not postgres %q: %d != %d", data, len(data), length)
-		return nil
+		return nil, 0
+	}
+
+	if dstPort == 5432 {
+		confidence = .1
 	}
 
 	major := int(binary.BigEndian.Uint16(data[4:]))
@@ -42,15 +46,19 @@ func detectPostgreSQLClient(dir Direction, data []byte) *Protocol {
 				Minor: minor,
 				Patch: -1,
 			},
-		}
+		}, confidence + .75
 	}
-	return nil
+	return nil, 0
 }
 
-func detectPostgreSQLServer(dir Direction, data []byte) *Protocol {
+func detectPostgreSQLServer(dir Direction, data []byte, srcPort, dstPort int) (proto *Protocol, confidence float64) {
 	// A server message needs at least 5 bytes (type + length).
 	if len(data) < 5 {
-		return nil
+		return nil, 0
+	}
+
+	if srcPort == 5432 {
+		confidence = .1
 	}
 
 	// All server messages (and subsequent client messages) are tagged with a single-byte type.
@@ -62,9 +70,9 @@ func detectPostgreSQLServer(dir Direction, data []byte) *Protocol {
 		'Z', // ReadyForQuery
 		'E', // ErrorResponse
 		'N': // NoticeResponse
-		return &Protocol{Name: ProtocolPostgreSQL}
+		return &Protocol{Name: ProtocolPostgreSQL}, confidence + .65
 
 	default:
-		return nil
+		return nil, 0
 	}
 }