bssh/
node.rs

1// Copyright 2025 Lablup Inc. and Jeongkyu Shin
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use anyhow::{Context, Result};
16use std::fmt;
17
18#[derive(Debug, Clone, PartialEq)]
19pub struct Node {
20    pub host: String,
21    pub port: u16,
22    pub username: String,
23}
24
25impl Node {
26    pub fn new(host: String, port: u16, username: String) -> Self {
27        Self {
28            host,
29            port,
30            username,
31        }
32    }
33
34    pub fn parse(node_str: &str, default_user: Option<&str>) -> Result<Self> {
35        // Parse formats:
36        // - host
37        // - host:port
38        // - user@host
39        // - user@host:port
40
41        let (user_part, host_part) = if let Some(at_pos) = node_str.find('@') {
42            let user = &node_str[..at_pos];
43            let rest = &node_str[at_pos + 1..];
44            (Some(user), rest)
45        } else {
46            (None, node_str)
47        };
48
49        let (host, port) = if let Some(colon_pos) = host_part.rfind(':') {
50            let host = &host_part[..colon_pos];
51            let port_str = &host_part[colon_pos + 1..];
52            let port = port_str.parse::<u16>().context("Invalid port number")?;
53            (host, port)
54        } else {
55            (host_part, 22)
56        };
57
58        let username = user_part
59            .or(default_user)
60            .map(|s| s.to_string())
61            .unwrap_or_else(|| {
62                std::env::var("USER")
63                    .or_else(|_| std::env::var("USERNAME"))
64                    .or_else(|_| std::env::var("LOGNAME"))
65                    .unwrap_or_else(|_| {
66                        // Try to get current user from system
67                        #[cfg(unix)]
68                        {
69                            whoami::username()
70                        }
71                        #[cfg(not(unix))]
72                        {
73                            "user".to_string()
74                        }
75                    })
76            });
77
78        Ok(Node {
79            host: host.to_string(),
80            port,
81            username,
82        })
83    }
84
85    pub fn address(&self) -> String {
86        format!("{}:{}", self.host, self.port)
87    }
88}
89
90impl fmt::Display for Node {
91    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
92        write!(f, "{}@{}:{}", self.username, self.host, self.port)
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99
100    #[test]
101    fn test_parse_host_only() {
102        let node = Node::parse("example.com", None).unwrap();
103        assert_eq!(node.host, "example.com");
104        assert_eq!(node.port, 22);
105    }
106
107    #[test]
108    fn test_parse_host_with_port() {
109        let node = Node::parse("example.com:2222", None).unwrap();
110        assert_eq!(node.host, "example.com");
111        assert_eq!(node.port, 2222);
112    }
113
114    #[test]
115    fn test_parse_user_and_host() {
116        let node = Node::parse("admin@example.com", None).unwrap();
117        assert_eq!(node.username, "admin");
118        assert_eq!(node.host, "example.com");
119        assert_eq!(node.port, 22);
120    }
121
122    #[test]
123    fn test_parse_full_format() {
124        let node = Node::parse("admin@example.com:2222", None).unwrap();
125        assert_eq!(node.username, "admin");
126        assert_eq!(node.host, "example.com");
127        assert_eq!(node.port, 2222);
128    }
129
130    #[test]
131    fn test_parse_with_default_user() {
132        let node = Node::parse("example.com", Some("default_user")).unwrap();
133        assert_eq!(node.username, "default_user");
134    }
135
136    #[test]
137    fn test_parse_uses_current_user_when_no_default() {
138        // When no user is specified, it should use current user from environment
139        let node = Node::parse("example.com", None).unwrap();
140        // Should not be "root" unless the current user is actually root
141        let current_user = std::env::var("USER")
142            .or_else(|_| std::env::var("USERNAME"))
143            .or_else(|_| std::env::var("LOGNAME"))
144            .unwrap_or_else(|_| whoami::username());
145        assert_eq!(node.username, current_user);
146        // Specifically verify it doesn't default to root when we're not root
147        if current_user != "root" {
148            assert_ne!(node.username, "root");
149        }
150    }
151}
bssh/node.rs

bssh/
node.rs
