Fix auth

py/sentry_relay/auth.py

@@ -0,0 +1,109 @@
+from ._relay_pyo3 import (
+    _generate_key_pair,
+    generate_relay_id,
+    _create_register_challenge,
+    _validate_register_response,
+    is_version_supported,
+    RustPublicKey,
+    RustSecretKey,
+    UnpackErrorBadSignature,
+)
+
+import json
+import uuid
+
+class SecretKey:
+    def __init__(self, inner):
+        self.inner = inner
+
+    @classmethod
+    def parse(cls, string):
+        return SecretKey(RustSecretKey.parse(string))
+
+    def sign(self, value):
+        return self.inner.sign(value)
+
+    def pack(self, data):
+        # TODO(@anonrig): Look into separators requirement
+        packed = json.dumps(data, separators=(",", ":")).encode()
+        return packed, self.sign(packed)
+
+    def __str__(self):
+        return self.inner.__str__()
+
+    def __repr__(self):
+        return self.inner.__repr__()
+
+
+class PublicKey:
+    def __init__(self, inner):
+        self.inner = inner
+
+    @classmethod
+    def parse(cls, string):
+        return SecretKey(RustPublicKey.parse(string))
+
+    def verify(self, buf, sig, max_age=None):
+        return self.inner.verify(buf, sig, max_age)
+
+    def unpack(
+        self,
+        buf,
+        sig,
+        max_age=None,
+    ):
+        if not self.verify(buf, sig, max_age):
+            raise UnpackErrorBadSignature()
+        return json.loads(buf)
+
+    def __str__(self):
+        return self.inner.__str__()
+
+    def __repr__(self):
+        return self.inner.__repr__()
+
+
+
+
+def generate_key_pair():
+    rsk, rpk = _generate_key_pair()
+    return (SecretKey(rsk), PublicKey(rpk))
+
+def create_register_challenge(
+    data,
+    signature,
+    secret,
+    max_age=60,
+    ):
+
+    challenge = _create_register_challenge(
+        data,
+        signature,
+        secret,
+        max_age,
+    )
+
+    return {
+        "relay_id": uuid.UUID(challenge["relay_id"]),
+        "token": challenge["token"],
+    }
+
+def validate_register_response(
+    data,
+    signature,
+    secret,
+    max_age=60,
+    ):
+    response = _validate_register_response(
+        data,
+        signature,
+        secret,
+        max_age,
+    )
+
+    return {
+        "relay_id": uuid.UUID(response["relay_id"]),
+        "token": response["token"],
+        "public_key": response["public_key"],
+        "version": response["version"],
+    }

py/sentry_relay/exceptions.py

@@ -0,0 +1,17 @@
+from ._relay_pyo3 import (
+        RelayError,
+        Unknown,
+        InvalidJsonError,
+        KeyParseErrorBadEncoding,
+        KeyParseErrorBadKey,
+        UnpackErrorBadSignature,
+        UnpackErrorBadPayload,
+        UnpackErrorSignatureExpired,
+        UnpackErrorBadEncoding,
+        ProcessingErrorInvalidTransaction,
+        ProcessingErrorInvalidGeoIp,
+        InvalidReleaseErrorTooLong,
+        InvalidReleaseErrorRestrictedName,
+        InvalidReleaseErrorBadCharacters,
+        RelayErrorCode
+    )

py/tests/test_auth.py

@@ -2,6 +2,8 @@
 import sentry_relay
 import pytest
 
+from sentry_relay import auth, exceptions
+
 
 UPSTREAM_SECRET = "secret"
 
@@ -18,19 +20,20 @@
 
 
 def test_basic_key_functions():
-    sk, pk = sentry_relay.generate_key_pair()
+    sk, pk = sentry_relay.auth.generate_key_pair()
+    print(pk)
     signature = sk.sign(b"some secret data")
     assert pk.verify(b"some secret data", signature)
     assert not pk.verify(b"some other data", signature)
 
     packed, signature = sk.pack({"foo": "bar"})
     pk.unpack(packed, signature)
-    with pytest.raises(sentry_relay.UnpackErrorBadSignature):
+    with pytest.raises(sentry_relay.exceptions.UnpackErrorBadSignature):
         pk.unpack(b"haha", signature)
 
 
 def test_challenge_response():
-    resp = sentry_relay.create_register_challenge(
+    resp = sentry_relay.auth.create_register_challenge(
         REQUEST,
         REQUEST_SIG,
         UPSTREAM_SECRET,
@@ -42,22 +45,22 @@ def test_challenge_response():
 
 
 def test_challenge_response_validation_errors():
-    with pytest.raises(sentry_relay.UnpackErrorSignatureExpired):
-        sentry_relay.create_register_challenge(
+    with pytest.raises(sentry_relay.exceptions.UnpackErrorSignatureExpired):
+        sentry_relay.auth.create_register_challenge(
             REQUEST,
             REQUEST_SIG,
             UPSTREAM_SECRET,
             max_age=1,
         )
-    with pytest.raises(sentry_relay.UnpackErrorBadPayload):
-        sentry_relay.create_register_challenge(
+    with pytest.raises(sentry_relay.exceptions.UnpackErrorBadPayload):
+        sentry_relay.auth.create_register_challenge(
             REQUEST + b"__broken",
             REQUEST_SIG,
             UPSTREAM_SECRET,
             max_age=0,
         )
-    with pytest.raises(sentry_relay.UnpackErrorBadSignature):
-        sentry_relay.create_register_challenge(
+    with pytest.raises(sentry_relay.exceptions.UnpackErrorBadSignature):
+        sentry_relay.auth.create_register_challenge(
             REQUEST,
             REQUEST_SIG + "__broken",
             UPSTREAM_SECRET,
@@ -66,7 +69,7 @@ def test_challenge_response_validation_errors():
 
 
 def test_register_response():
-    resp = sentry_relay.validate_register_response(
+    resp = sentry_relay.auth.validate_register_response(
         RESPONSE,
         RESPONSE_SIG,
         UPSTREAM_SECRET,
@@ -78,8 +81,8 @@ def test_register_response():
 
 
 def test_is_version_supported():
-    assert sentry_relay.is_version_supported("99.99.99")
+    assert sentry_relay.auth.is_version_supported("99.99.99")
 
     # These can be updated when deprecating legacy versions:
-    assert sentry_relay.is_version_supported("")
-    assert sentry_relay.is_version_supported(None)
+    assert sentry_relay.auth.is_version_supported("")
+    assert sentry_relay.auth.is_version_supported(None)

py/tests/test_processing.py

@@ -1,5 +1,5 @@
 import sentry_relay
-from sentry_relay import processing
+from sentry_relay import processing, exceptions
 
 import pytest
 
@@ -219,7 +219,7 @@ def test_parse_release():
 
 
 def test_parse_release_error():
-    with pytest.raises(sentry_relay.errors.InvalidReleaseErrorBadCharacters):
+    with pytest.raises(sentry_relay.exceptions.InvalidReleaseErrorBadCharacters):
         sentry_relay.processing.parse_release("/var/foo/foo")
 
 

relay-pii/src/generate_selectors.rs

@@ -1,6 +1,5 @@
 use std::collections::BTreeSet;
 
-use pyo3::prelude::*;
 use relay_event_schema::processor::{
     self, Pii, ProcessValue, ProcessingResult, ProcessingState, Processor, ValueType,
 };

relay-pyo3/src/auth.rs

@@ -9,9 +9,10 @@ use relay_auth::{
     KeyParseError, PublicKey, RegisterRequest, RegisterResponse, RelayVersion, SecretKey,
 };
 
+use crate::exceptions::pyerr_from_unpack_error;
 use crate::utils::extract_bytes_or_str;
 
-#[pyclass(name = "PublicKey")]
+#[pyclass(name = "RustPublicKey")]
 #[derive(Clone, Eq, PartialEq)]
 pub struct PyPublicKey(PublicKey);
 
@@ -40,20 +41,6 @@ impl PyPublicKey {
         }
     }
 
-    fn unpack(
-        &self,
-        buf: &Bound<PyBytes>,
-        sig: &Bound<PyAny>,
-        max_age: Option<Duration>,
-    ) -> PyResult<()> {
-        if !self.verify(buf, sig, max_age)? {
-            todo!("Return error");
-        }
-
-        let _buf = buf.as_bytes();
-        todo!("Deserialize into dict");
-    }
-
     fn __str__(&self) -> String {
         self.0.to_string()
     }
@@ -64,7 +51,7 @@ impl PyPublicKey {
     }
 }
 
-#[pyclass(name = "SecretKey")]
+#[pyclass(name = "RustSecretKey")]
 #[derive(Clone, Eq, PartialEq)]
 pub struct PySecretKey(SecretKey);
 
@@ -83,8 +70,6 @@ impl PySecretKey {
         self.0.sign(value.as_bytes())
     }
 
-    // TODO: implement pack on the Python side
-
     fn __str__(&self) -> String {
         self.0.to_string()
     }
@@ -95,7 +80,7 @@ impl PySecretKey {
     }
 }
 
-#[pyfunction]
+#[pyfunction(name = "_generate_key_pair")]
 fn generate_key_pair() -> (PySecretKey, PyPublicKey) {
     let (secret, public) = relay_auth::generate_key_pair();
     (PySecretKey(secret), PyPublicKey(public))
@@ -108,46 +93,45 @@ fn generate_relay_id() -> [u8; 16] {
     id.into_bytes()
 }
 
-#[pyfunction]
+#[pyfunction(name = "_create_register_challenge")]
 #[pyo3(signature = (data, signature, secret, max_age = 60))]
 fn create_register_challenge(
     data: &Bound<PyBytes>,
     signature: &Bound<PyAny>,
     secret: &Bound<PyAny>,
-    max_age: i64,
+    max_age: u32,
 ) -> PyResult<HashMap<&'static str, String>> {
-    let max_age = Duration::seconds(max_age);
+    let max_age = (max_age > 0).then_some(Duration::seconds(max_age.into()));
     let signature = extract_bytes_or_str(signature)?;
     let secret = extract_bytes_or_str(secret)?;
 
-    let req = RegisterRequest::bootstrap_unpack(data.as_bytes(), signature, Some(max_age))
-        .map_err(|e| PyValueError::new_err(e.to_string()))?;
+    let req = RegisterRequest::bootstrap_unpack(data.as_bytes(), signature, max_age)
+        .map_err(pyerr_from_unpack_error)?;
 
     let challenge = req.into_challenge(secret.as_bytes());
 
     let mut output = HashMap::new();
 
-    // TODO: Parse the UUID on the Python side
     output.insert("relay_id", challenge.relay_id().to_string());
     output.insert("token", challenge.token().to_owned());
 
     Ok(output)
 }
 
-#[pyfunction]
+#[pyfunction(name = "_validate_register_response")]
 #[pyo3(signature = (data, signature, secret, max_age = 60))]
 fn validate_register_response(
     data: &Bound<PyBytes>,
     signature: &Bound<PyAny>,
     secret: &Bound<PyAny>,
-    max_age: i64,
+    max_age: u32,
 ) -> PyResult<HashMap<&'static str, String>> {
-    let max_age = Duration::seconds(max_age);
+    let max_age = (max_age > 0).then_some(Duration::seconds(max_age.into()));
     let signature = extract_bytes_or_str(signature)?;
     let secret = extract_bytes_or_str(secret)?;
 
     let (response, state) =
-        RegisterResponse::unpack(data.as_bytes(), signature, secret.as_bytes(), Some(max_age))
+        RegisterResponse::unpack(data.as_bytes(), signature, secret.as_bytes(), max_age)
             .map_err(|e| PyValueError::new_err(e.to_string()))?;
 
     let mut output = HashMap::new();

relay-pyo3/src/exceptions.rs

@@ -150,6 +150,15 @@ make_error!(InvalidReleaseErrorTooLong);
 make_error!(InvalidReleaseErrorRestrictedName);
 make_error!(InvalidReleaseErrorBadCharacters);
 
+pub(crate) fn pyerr_from_unpack_error(value: UnpackError) -> PyErr {
+    match value {
+        UnpackError::BadSignature => UnpackErrorBadSignature::new().into(),
+        UnpackError::BadEncoding => UnpackErrorBadEncoding::new().into(),
+        UnpackError::BadPayload(_) => UnpackErrorBadPayload::new().into(),
+        UnpackError::SignatureExpired => UnpackErrorSignatureExpired::new().into(),
+    }
+}
+
 #[pymodule]
 pub fn exceptions(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_class::<RelayError>()?;

relay-pyo3/src/lib.rs

@@ -1,15 +1,16 @@
 use pyo3::prelude::*;
-use pyo3::wrap_pymodule;
 
 mod auth;
 mod codeowners;
 mod consts;
-mod exceptions;
+pub(crate) mod exceptions;
 mod processing;
 mod utils;
 
 #[pymodule]
 fn _relay_pyo3(m: &Bound<PyModule>) -> PyResult<()> {
+    auth::auth(m)?;
+    exceptions::exceptions(m)?;
     processing::processing(m)?;
 
     Ok(())