Line data    Source code

       1              : use lazy_static::lazy_static;
       2              : use regex_automata::util::captures::Captures;
       3              : use regex_automata::meta::Regex;
       4              : 
       5              : #[derive(Debug, PartialEq)]
       6              : pub enum GrpcProxyRoute {
       7              :     Channels,
       8              :     Channel { channel_id: usize } ,
       9              :     ChannelUnary { channel_id: usize },
      10              :     ChannelStreams { channel_id: usize },
      11              :     ChannelStream { channel_id: usize, stream_id: usize },
      12              : }
      13              : 
      14              : lazy_static! {
      15              :     static ref ROUTES: Regex = Regex::new_many(&[
      16              :         r"^/grpc/channels$",
      17              :         r"^/grpc/channel/(\d+)$",
      18              :         r"^/grpc/channel/(\d+)/unary$",
      19              :         r"^/grpc/channel/(\d+)/streams$",
      20              :         r"^/grpc/channel/(\d+)/stream/(\d+)$",
      21              :     ]).unwrap();
      22              : }
      23              : 
      24           65 : pub fn parse_grpc_proxy_path(path: &str) -> Option<GrpcProxyRoute> {
      25           65 :     let mut all = Captures::all(ROUTES.group_info().clone());
      26           65 :     ROUTES.captures(path, &mut all);
      27           65 :     match all.pattern().map(|p| p.as_usize()) {
      28              :         Some(0) => {
      29           21 :             Some(GrpcProxyRoute::Channels)
      30              :         }
      31              :         Some(1) => {
      32            2 :             let channel_id = path[all.get_group(1).unwrap()].parse().unwrap_or_default();
      33            2 :             Some(GrpcProxyRoute::Channel { channel_id })
      34              :         }
      35              :         Some(2) => {
      36            9 :             let channel_id = path[all.get_group(1).unwrap()].parse().unwrap_or_default();
      37            9 :             Some(GrpcProxyRoute::ChannelUnary { channel_id })
      38              :         }
      39              :         Some(3) => {
      40            9 :             let channel_id = path[all.get_group(1).unwrap()].parse().unwrap_or_default();
      41            9 :             Some(GrpcProxyRoute::ChannelStreams { channel_id })
      42              :         }
      43              :         Some(4) => {
      44           16 :             let channel_id = path[all.get_group(1).unwrap()].parse().unwrap_or_default();
      45           16 :             let stream_id = path[all.get_group(2).unwrap()].parse().unwrap_or_default();
      46           16 :             Some(GrpcProxyRoute::ChannelStream { channel_id, stream_id })
      47              :         }
      48              :         _ => {
      49            8 :             None
      50              :         }
      51              :     }
      52           65 : }
      53              : 
      54              : #[cfg(test)]
      55              : mod test {
      56              :     use super::*;
      57              :     use anyhow::Result;
      58              : 
      59              :     #[tokio::test]
      60            1 :     async fn test_valid_routes() -> Result<()> {
      61            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channels"), Some(GrpcProxyRoute::Channels));
      62            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/1"), Some(GrpcProxyRoute::Channel { channel_id: 1 }));
      63            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/1/unary"), Some(GrpcProxyRoute::ChannelUnary { channel_id: 1 }));
      64            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/1/streams"), Some(GrpcProxyRoute::ChannelStreams { channel_id: 1 }));
      65            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/1/stream/2"), Some(GrpcProxyRoute::ChannelStream { channel_id: 1, stream_id: 2 }));
      66            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/123/stream/456"), Some(GrpcProxyRoute::ChannelStream { channel_id: 123, stream_id: 456 }));
      67            2 :         Ok(())
      68            1 :     }
      69              : 
      70              :     #[tokio::test]
      71            1 :     async fn test_invalid_routes() -> Result<()> {
      72            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/foo"), None);
      73            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channels/foo"), None);
      74            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/foo"), None);
      75            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/1/foo"), None);
      76            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/1/stream/foo"), None);
      77            1 :         assert_eq!(parse_grpc_proxy_path("/grpc/channel/1/stream/2/foo"), None);
      78            2 :         Ok(())
      79            1 :     }
      80              : }