feat(flags): Do token validation and extract distinct id

feature-flags/Cargo.toml

@@ -24,6 +24,7 @@ redis = { version = "0.23.3", features = [
 serde = { workspace = true }
 serde_json = { workspace = true }
 thiserror = { workspace = true }
+serde-pickle = { version = "1.1.1"}
 
 [lints]
 workspace = true

feature-flags/src/api.rs

@@ -25,6 +25,9 @@ pub enum FlagError {
     #[error("failed to parse request: {0}")]
     RequestParsingError(#[from] serde_json::Error),
 
+    #[error("failed to parse redis data: {0}")]
+    DataParsingError(#[from] serde_pickle::Error),
+
     #[error("Empty distinct_id in request")]
     EmptyDistinctId,
     #[error("No distinct_id in request")]
@@ -44,6 +47,7 @@ impl IntoResponse for FlagError {
         match self {
             FlagError::RequestDecodingError(_)
             | FlagError::RequestParsingError(_)
+            | FlagError::DataParsingError(_)
             | FlagError::EmptyDistinctId
             | FlagError::MissingDistinctId => (StatusCode::BAD_REQUEST, self.to_string()),
 

feature-flags/src/config.rs

@@ -4,7 +4,7 @@ use envconfig::Envconfig;
 
 #[derive(Envconfig, Clone)]
 pub struct Config {
-    #[envconfig(default = "127.0.0.1:0")]
+    #[envconfig(default = "127.0.0.1:3001")]
     pub address: SocketAddr,
 
     #[envconfig(default = "postgres://posthog:posthog@localhost:15432/test_database")]

feature-flags/src/lib.rs

@@ -5,3 +5,4 @@ pub mod router;
 pub mod server;
 pub mod v0_endpoint;
 pub mod v0_request;
+pub mod team;

feature-flags/src/redis.rs

@@ -2,20 +2,23 @@ use std::time::Duration;
 
 use anyhow::Result;
 use async_trait::async_trait;
-use redis::AsyncCommands;
+use redis::{AsyncCommands, RedisError};
 use tokio::time::timeout;
 
 // average for all commands is <10ms, check grafana
 const REDIS_TIMEOUT_MILLISECS: u64 = 10;
 
 /// A simple redis wrapper
 /// Copied from capture/src/redis.rs.
-/// TODO: Modify this to support hincrby, get, and set commands.
+/// TODO: Modify this to support hincrby
 
 #[async_trait]
 pub trait Client {
     // A very simplified wrapper, but works for our usage
     async fn zrangebyscore(&self, k: String, min: String, max: String) -> Result<Vec<String>>;
+
+    async fn get(&self, k: String) -> Result<String>;
+    async fn set(&self, k: String, v: String) -> Result<()>;
 }
 
 pub struct RedisClient {
@@ -40,38 +43,31 @@ impl Client for RedisClient {
 
         Ok(fut?)
     }
-}
 
-// TODO: Find if there's a better way around this.
-#[derive(Clone)]
-pub struct MockRedisClient {
-    zrangebyscore_ret: Vec<String>,
-}
+    async fn get(&self, k: String) -> Result<String> {
+        let mut conn = self.client.get_async_connection().await?;
 
-impl MockRedisClient {
-    pub fn new() -> MockRedisClient {
-        MockRedisClient {
-            zrangebyscore_ret: Vec::new(),
-        }
-    }
+        let results = conn.get(k.clone());
+        // TODO: Is this safe? Should we be doing something else for error handling here?
+        let fut: Result<Vec<u8>, RedisError> = timeout(Duration::from_secs(REDIS_TIMEOUT_MILLISECS), results).await?;
 
-    pub fn zrangebyscore_ret(&mut self, ret: Vec<String>) -> Self {
-        self.zrangebyscore_ret = ret;
+        // TRICKY: We serialise data to json, then django pickles it.
+        // Here we deserialize the bytes using serde_pickle, to get the json string.
+        let string_response: String = serde_pickle::from_slice(&fut?, Default::default())?;
 
-        self.clone()
+        Ok(string_response)
     }
-}
 
-impl Default for MockRedisClient {
-    fn default() -> Self {
-        Self::new()
-    }
-}
+    async fn set(&self, k: String, v: String) -> Result<()> {
+        // TRICKY: We serialise data to json, then django pickles it.
+        // Here we serialize the json string to bytes using serde_pickle.
+        let bytes = serde_pickle::to_vec(&v, Default::default())?;
 
-#[async_trait]
-impl Client for MockRedisClient {
-    // A very simplified wrapper, but works for our usage
-    async fn zrangebyscore(&self, _k: String, _min: String, _max: String) -> Result<Vec<String>> {
-        Ok(self.zrangebyscore_ret.clone())
+        let mut conn = self.client.get_async_connection().await?;
+
+        let results = conn.set(k, bytes);
+        let fut = timeout(Duration::from_secs(REDIS_TIMEOUT_MILLISECS), results).await?;
+
+        Ok(fut?)
     }
-}
+}

feature-flags/src/v0_endpoint.rs

@@ -33,7 +33,7 @@ use crate::{
 )]
 #[debug_handler]
 pub async fn flags(
-    _state: State<router::State>,
+    state: State<router::State>,
     InsecureClientIp(ip): InsecureClientIp,
     meta: Query<FlagsQueryParams>,
     headers: HeaderMap,
@@ -59,19 +59,19 @@ pub async fn flags(
         .get("content-type")
         .map_or("", |v| v.to_str().unwrap_or(""))
     {
-        "application/x-www-form-urlencoded" => {
-            return Err(FlagError::RequestDecodingError(String::from(
-                "invalid form data",
-            )));
+        "application/json" => {
+            tracing::Span::current().record("content_type", "application/json");
+            FlagRequest::from_bytes(body)
         }
         ct => {
-            tracing::Span::current().record("content_type", ct);
-
-            FlagRequest::from_bytes(body)
+            return Err(FlagError::RequestDecodingError(format!(
+                "unsupported content type: {}",
+                ct
+            )));
         }
     }?;
 
-    let token = request.extract_and_verify_token()?;
+    let token = request.extract_and_verify_token(state.redis.clone()).await?;
 
     tracing::Span::current().record("token", &token);
 

feature-flags/src/v0_request.rs

@@ -1,11 +1,11 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 use bytes::Bytes;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use tracing::instrument;
 
-use crate::api::FlagError;
+use crate::{api::FlagError, redis::Client, team::Team};
 
 #[derive(Deserialize, Default)]
 pub struct FlagsQueryParams {
@@ -54,15 +54,18 @@ impl FlagRequest {
         Ok(serde_json::from_str::<FlagRequest>(&payload)?)
     }
 
-    pub fn extract_and_verify_token(&self) -> Result<String, FlagError> {
+    pub async fn extract_and_verify_token(&self, redis_client: Arc<dyn Client + Send + Sync>) -> Result<String, FlagError> {
         let token = match self {
             FlagRequest {
                 token: Some(token), ..
             } => token.to_string(),
             _ => return Err(FlagError::NoTokenError),
         };
-        // TODO: Get tokens from redis, confirm this one is valid
-        // validate_token(&token)?;
+
+        let team = Team::from_redis(redis_client, token.clone()).await?;
+
+        // TODO: Remove this, is useless, doing just for now because 
+        tracing::Span::current().record("team_id", &team.id);
         Ok(token)
     }
 }