Merge pull request #557 from benjamin-747/main

[mono] added pubkey verify in ssh connection

Cargo.toml

@@ -36,18 +36,18 @@ mono = { path = "mono" }
 
 anyhow = "1.0.86"
 serde = "1.0.205"
-serde_json = "1.0.122"
+serde_json = "1.0.128"
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"
 tracing-appender = "0.2"
 thiserror = "1.0.63"
 rand = "0.8.5"
 smallvec = "1.13.2"
-tokio = "1.39.2"
+tokio = "1.40.0"
 tokio-stream = "0.1.15"
 tokio-test = "0.4.4"
-clap = "4.5.14"
-async-trait = "0.1.81"
+clap = "4.5.17"
+async-trait = "0.1.82"
 async-stream = "0.3.5"
 bytes = "1.7.1"
 memchr = "2.7.4"

ceres/src/protocol/smart.rs

@@ -8,11 +8,9 @@ use tokio_stream::wrappers::ReceiverStream;
 use callisto::db_enums::RefType;
 use common::errors::ProtocolError;
 
-use crate::protocol::ZERO_ID;
-use crate::protocol::{
-    Capability, ServiceType, SideBind, SmartProtocol, TransportProtocol,
-};
 use crate::protocol::import_refs::RefCommand;
+use crate::protocol::ZERO_ID;
+use crate::protocol::{Capability, ServiceType, SideBind, SmartProtocol, TransportProtocol};
 
 const LF: char = '\n';
 
@@ -22,6 +20,7 @@ const NUL: char = '\0';
 
 pub const PKT_LINE_END_MARKER: &[u8; 4] = b"0000";
 
+// see https://git-scm.com/docs/protocol-capabilities
 // The atomic, report-status, report-status-v2, delete-refs, quiet,
 // and push-cert capabilities are sent and recognized by the receive-pack (push to server) process.
 const RECEIVE_CAP_LIST: &str = "report-status report-status-v2 delete-refs quiet atomic no-thin ";
@@ -31,8 +30,7 @@ const RECEIVE_CAP_LIST: &str = "report-status report-status-v2 delete-refs quiet
 const COMMON_CAP_LIST: &str = "side-band-64k ofs-delta agent=mega/0.1.0";
 
 // All other capabilities are only recognized by the upload-pack (fetch from server) process.
-const UPLOAD_CAP_LIST: &str =
-    "shallow deepen-since deepen-not deepen-relative multi_ack_detailed no-done include-tag ";
+const UPLOAD_CAP_LIST: &str = "multi_ack_detailed no-done include-tag ";
 
 impl SmartProtocol {
     /// # Retrieves the information about Git references (refs) for the specified service type.
@@ -401,8 +399,8 @@ pub fn read_pkt_line(bytes: &mut Bytes) -> (usize, Bytes) {
 pub mod test {
     use bytes::{Bytes, BytesMut};
     use callisto::db_enums::RefType;
-    use crate::protocol::import_refs::{CommandType, RefCommand};
 
+    use crate::protocol::import_refs::{CommandType, RefCommand};
     use crate::protocol::smart::{add_pkt_line_string, read_pkt_line, read_until_white_space};
     use crate::protocol::{Capability, SmartProtocol};
 

jupiter/callisto/src/ssh_keys.rs

@@ -10,6 +10,8 @@ pub struct Model {
     pub user_id: i64,
     #[sea_orm(column_type = "Text")]
     pub ssh_key: String,
+    #[sea_orm(column_type = "Text")]
+    pub finger: String,
     pub created_at: DateTime,
 }
 

jupiter/src/storage/user_storage.rs

@@ -41,11 +41,17 @@ impl UserStorage {
         Ok(())
     }
 
-    pub async fn save_ssh_key(&self, user_id: i64, ssh_key: &str) -> Result<(), MegaError> {
+    pub async fn save_ssh_key(
+        &self,
+        user_id: i64,
+        ssh_key: &str,
+        finger: &str,
+    ) -> Result<(), MegaError> {
         let model = ssh_keys::Model {
             id: generate_id(),
             user_id,
             ssh_key: ssh_key.to_owned(),
+            finger: finger.to_owned(),
             created_at: chrono::Utc::now().naive_utc(),
         };
         let a_model = model.into_active_model();
@@ -67,4 +73,15 @@ impl UserStorage {
             .await?;
         Ok(())
     }
+
+    pub async fn search_ssh_key_finger(
+        &self,
+        finger_print: &str,
+    ) -> Result<Vec<ssh_keys::Model>, MegaError> {
+        let res = ssh_keys::Entity::find()
+            .filter(ssh_keys::Column::Finger.eq(finger_print))
+            .all(self.get_connection())
+            .await?;
+        Ok(res)
+    }
 }

mono/src/api/user/user_router.rs

@@ -4,6 +4,7 @@ use axum::{
     routing::{get, post},
     Json, Router,
 };
+use russh_keys::parse_public_key_base64;
 
 use common::model::CommonResult;
 
@@ -32,11 +33,20 @@ async fn add_key(
     state: State<MonoApiServiceState>,
     Json(json): Json<AddSSHKey>,
 ) -> Result<Json<CommonResult<String>>, (StatusCode, String)> {
+    let key_data = json
+        .ssh_key
+        .split_whitespace()
+        .nth(1)
+        .ok_or("Invalid key format")
+        .unwrap();
+
+    let key = parse_public_key_base64(key_data).unwrap();
+
     let res = state
         .context
         .services
         .user_storage
-        .save_ssh_key(user.user_id, &json.ssh_key)
+        .save_ssh_key(user.user_id, &json.ssh_key, &key.fingerprint())
         .await;
     let res = match res {
         Ok(_) => CommonResult::success(None),

mono/src/git_protocol/http.rs

@@ -160,7 +160,7 @@ pub async fn git_receive_pack(
 }
 
 // Function to find the subsequence in a slice
-fn search_subsequence(chunk: &[u8], search: &[u8]) -> Option<usize> {
+pub fn search_subsequence(chunk: &[u8], search: &[u8]) -> Option<usize> {
     chunk.windows(search.len()).position(|s| s == search)
 }
 

mono/src/git_protocol/ssh.rs

@@ -1,14 +1,14 @@
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::str::FromStr;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 
 use async_trait::async_trait;
 use bytes::{Bytes, BytesMut};
 use chrono::{DateTime, Duration, Utc};
 use futures::{stream, StreamExt};
-use russh::server::{self, Auth, Msg, Response, Session};
-use russh::{Channel, ChannelId};
+use russh::server::{self, Auth, Msg, Session};
+use russh::{Channel, ChannelId, MethodSet};
 use russh_keys::key;
 use tokio::io::AsyncReadExt;
 
@@ -17,6 +17,9 @@ use ceres::protocol::smart::{self};
 use ceres::protocol::ServiceType;
 use ceres::protocol::{SmartProtocol, TransportProtocol};
 use jupiter::context::Context;
+use tokio::sync::Mutex;
+
+use crate::git_protocol::http::search_subsequence;
 
 type ClientMap = HashMap<(usize, ChannelId), Channel<Msg>>;
 #[allow(dead_code)]
@@ -26,9 +29,8 @@ pub struct SshServer {
     pub clients: Arc<Mutex<ClientMap>>,
     pub id: usize,
     pub context: Context,
-    // TODO: consider is it a good choice to bind data here, find a better solution to bind data with ssh client
     pub smart_protocol: Option<SmartProtocol>,
-    pub data_combined: Vec<u8>,
+    pub data_combined: BytesMut,
 }
 
 impl server::Server for SshServer {
@@ -51,7 +53,7 @@ impl server::Handler for SshServer {
     ) -> Result<bool, Self::Error> {
         tracing::info!("SshServer::channel_open_session:{}", channel.id());
         {
-            let mut clients = self.clients.lock().unwrap();
+            let mut clients = self.clients.lock().await;
             clients.insert((self.id, channel.id()), channel);
         }
         Ok(true)
@@ -129,25 +131,24 @@ impl server::Handler for SshServer {
         user: &str,
         public_key: &key::PublicKey,
     ) -> Result<Auth, Self::Error> {
-        tracing::info!("auth_publickey: {} / {:?}", user, public_key);
-        Ok(Auth::Accept)
-    }
-
-    async fn auth_keyboard_interactive(
-        &mut self,
-        _: &str,
-        _: &str,
-        _: Option<Response<'async_trait>>,
-    ) -> Result<Auth, Self::Error> {
-        tracing::info!("auth_keyboard_interactive");
-        Ok(Auth::Accept)
-    }
-
-    // TODO! disable password auth
-    async fn auth_password(&mut self, user: &str, password: &str) -> Result<Auth, Self::Error> {
-        tracing::info!("auth_password: {} / {}", user, password);
-        // in this example implementation, any username/password combination is accepted
-        Ok(Auth::Accept)
+        tracing::info!(
+            "auth_publickey: {} / {:?}/ {}",
+            user,
+            public_key.name(),
+            public_key.fingerprint()
+        );
+        let fingerprint = public_key.fingerprint();
+        let stg = self.context.services.user_storage.clone();
+        let res = stg.search_ssh_key_finger(&fingerprint).await.unwrap();
+        if !res.is_empty() {
+            tracing::info!("Client public key verified successfully!");
+            Ok(Auth::Accept)
+        } else {
+            tracing::warn!("Client public key verification failed!");
+            Ok(Auth::Reject {
+                proceed_with_methods: Some(MethodSet::PUBLICKEY),
+            })
+        }
     }
 
     async fn data(
@@ -168,7 +169,7 @@ impl server::Handler for SshServer {
                 self.handle_upload_pack(channel, data, session).await;
             }
             ServiceType::ReceivePack => {
-                self.data_combined.extend(data);
+                self.data_combined.extend_from_slice(data);
             }
         };
         session.channel_success(channel);
@@ -187,7 +188,7 @@ impl server::Handler for SshServer {
         }
 
         {
-            let mut clients = self.clients.lock().unwrap();
+            let mut clients = self.clients.lock().await;
             clients.remove(&(self.id, channel));
         }
         session.exit_status_request(channel, 0000);
@@ -226,16 +227,27 @@ impl SshServer {
 
     async fn handle_receive_pack(&mut self, channel: ChannelId, session: &mut Session) {
         let smart_protocol = self.smart_protocol.as_mut().unwrap();
+        let data = self.data_combined.split().freeze();
+        let mut data_stream = Box::pin(stream::once(async move { Ok(data) }));
+        let mut report_status = Bytes::new();
+
+        while let Some(chunk) = data_stream.next().await {
+            let chunk = chunk.unwrap();
+
+            if let Some(pos) = search_subsequence(&chunk, b"PACK") {
+                smart_protocol.git_receive_pack_protocol(Bytes::copy_from_slice(&chunk[..pos]));
+                let remaining_bytes = Bytes::copy_from_slice(&chunk[pos..]);
+                let remaining_stream =
+                    stream::once(async { Ok(remaining_bytes) }).chain(data_stream);
+                report_status = smart_protocol
+                    .git_receive_pack_stream(Box::pin(remaining_stream))
+                    .await
+                    .unwrap();
+                break;
+            }
+        }
 
-        let data = self.data_combined.clone();
-        let stream = stream::once(async move {
-            Ok(Bytes::from(data))
-        });
-        let buf = smart_protocol
-            .git_receive_pack_stream(Box::pin(stream))
-            .await
-            .unwrap();
-        tracing::info!("report status: {:?}", buf);
-        session.data(channel, buf.to_vec().into());
+        tracing::info!("report status: {:?}", report_status);
+        session.data(channel, report_status.to_vec().into());
     }
 }

mono/src/server/ssh_server.rs

@@ -4,9 +4,10 @@ use std::io::Read;
 use std::net::SocketAddr;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 
 use anyhow::Result;
+use bytes::BytesMut;
 use clap::Args;
 
 use common::config::Config;
@@ -18,6 +19,7 @@ use russh_keys::key::KeyPair;
 
 use common::model::CommonOptions;
 use jupiter::context::Context;
+use tokio::sync::Mutex;
 
 use crate::git_protocol::ssh::SshServer;
 
@@ -32,7 +34,7 @@ pub struct SshOptions {
 
 #[derive(Args, Clone, Debug)]
 pub struct SshCustom {
-    #[arg(long, default_value_t = 8100)]
+    #[arg(long, default_value_t = 2222)]
     ssh_port: u16,
 
     #[arg(long, value_name = "FILE")]
@@ -77,7 +79,7 @@ pub async fn start_server(config: Config, command: &SshOptions) {
         id: 0,
         context,
         smart_protocol: None,
-        data_combined: Vec::new(),
+        data_combined: BytesMut::new(),
     };
     let server_url = format!("{}:{}", host, ssh_port);
     let addr = SocketAddr::from_str(&server_url).unwrap();

sql/postgres/pg_20240905__init.sql

@@ -292,7 +292,8 @@ CREATE TABLE IF NOT EXISTS "ssh_keys" (
   "id" BIGINT PRIMARY KEY,
   "user_id" BIGINT NOT NULL,
   "ssh_key" TEXT NOT NULL,
+  "finger" TEXT NOT NULL,
   "created_at" TIMESTAMP NOT NULL
 );
 CREATE INDEX "idx_user_id" ON "ssh_keys" ("user_id");
-CREATE INDEX "idx_ssh_key_expression" ON "ssh_keys"  ((left(ssh_key, 32)));
+CREATE INDEX "idx_ssh_key_finger" ON "ssh_keys" ((left(finger, 8)));

sql/sqlite/sqlite_20240905_init.sql

@@ -292,7 +292,8 @@ CREATE TABLE IF NOT EXISTS "ssh_keys" (
   "id" BIGINT PRIMARY KEY,
   "user_id" BIGINT NOT NULL,
   "ssh_key" TEXT NOT NULL,
+  "finger" TEXT NOT NULL,
   "created_at" TIMESTAMP NOT NULL
 );
 CREATE INDEX "idx_user_id" ON "ssh_keys" ("user_id");
-CREATE INDEX "idx_ssh_key_expression" ON "ssh_keys" ("ssh_key");
+CREATE INDEX "idx_ssh_key_finger" ON "ssh_keys" ("finger");