Skip to content

Commit bbd8205

Browse files
baloowiktor-k
baloo
authored and
wiktor-k
committed
proto: rewrite error handling, remove a todo
Signed-off-by: Arthur Gautier <[email protected]>
1 parent 003513a commit bbd8205

File tree

4 files changed

+41
-38
lines changed

4 files changed

+41
-38
lines changed

‎src/agent.rs

+3-14
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ use std::marker::Unpin;
1818
use std::mem::size_of;
1919

2020
use super::error::AgentError;
21-
use super::proto::message::Message;
21+
use super::proto::{message::Message, ProtoError};
2222

2323
#[derive(Debug)]
2424
pub struct MessageCodec;
@@ -40,11 +40,6 @@ impl Decoder for MessageCodec {
4040
return Ok(None);
4141
}
4242

43-
//use std::io::Write;
44-
//let mut file = std::fs::File::create(uuid::Uuid::new_v4().to_string())?;
45-
//file.write_all(bytes)?;
46-
//drop(file);
47-
4843
let message: Message = Message::decode(&mut bytes)?;
4944
src.advance(size_of::<u32>() + length);
5045
Ok(Some(message))
@@ -58,16 +53,10 @@ impl Encoder<Message> for MessageCodec {
5853
let mut bytes = Vec::new();
5954

6055
let len = item.encoded_len().unwrap() as u32;
61-
len.encode(&mut bytes)?;
56+
len.encode(&mut bytes).map_err(ProtoError::SshEncoding)?;
6257

63-
item.encode(&mut bytes)?;
58+
item.encode(&mut bytes).map_err(ProtoError::SshEncoding)?;
6459
dst.put(&*bytes);
65-
//use std::io::Write;
66-
//let mut file = std::fs::File::create(uuid::Uuid::new_v4().to_string())?;
67-
//let mut bytes = Vec::new();
68-
//item.encode(&mut bytes)?;
69-
//file.write_all(&bytes)?;
70-
//drop(file);
7160

7261
Ok(())
7362
}

‎src/error.rs

+12-10
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,25 @@
11
use std::io;
22

3+
use crate::proto::ProtoError;
4+
35
#[derive(Debug)]
46
pub enum AgentError {
5-
Ssh(ssh_key::Error),
6-
Proto(ssh_encoding::Error),
7+
//Ssh(ssh_key::Error),
8+
Proto(ProtoError),
79
IO(io::Error),
810
}
911

10-
impl From<ssh_encoding::Error> for AgentError {
11-
fn from(e: ssh_encoding::Error) -> AgentError {
12+
impl From<ProtoError> for AgentError {
13+
fn from(e: ProtoError) -> AgentError {
1214
AgentError::Proto(e)
1315
}
1416
}
1517

16-
impl From<ssh_key::Error> for AgentError {
17-
fn from(e: ssh_key::Error) -> AgentError {
18-
AgentError::Ssh(e)
19-
}
20-
}
18+
//impl From<ssh_key::Error> for AgentError {
19+
// fn from(e: ssh_key::Error) -> AgentError {
20+
// AgentError::Ssh(e)
21+
// }
22+
//}
2123

2224
impl From<io::Error> for AgentError {
2325
fn from(e: io::Error) -> AgentError {
@@ -28,7 +30,7 @@ impl From<io::Error> for AgentError {
2830
impl std::fmt::Display for AgentError {
2931
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3032
match self {
31-
AgentError::Ssh(e) => write!(f, "Agent: Ssh key error: {e}"),
33+
//AgentError::Ssh(e) => write!(f, "Agent: Ssh key error: {e}"),
3234
AgentError::Proto(proto) => write!(f, "Agent: Protocol error: {}", proto),
3335
AgentError::IO(error) => write!(f, "Agent: I/O error: {}", error),
3436
}

‎src/proto/error.rs

+5
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ pub enum ProtoError {
99
IO(io::Error),
1010
SshEncoding(ssh_encoding::Error),
1111
SshKey(ssh_key::Error),
12+
UnsupportedCommand { command: u8 },
1213
}
1314

1415
impl From<ProtoError> for () {
@@ -48,6 +49,7 @@ impl std::error::Error for ProtoError {
4849
ProtoError::IO(e) => Some(e),
4950
ProtoError::SshEncoding(e) => Some(e),
5051
ProtoError::SshKey(e) => Some(e),
52+
ProtoError::UnsupportedCommand { .. } => None,
5153
}
5254
}
5355
}
@@ -61,6 +63,9 @@ impl std::fmt::Display for ProtoError {
6163
ProtoError::IO(_) => f.write_str("I/O Error"),
6264
ProtoError::SshEncoding(_) => f.write_str("SSH encoding Error"),
6365
ProtoError::SshKey(e) => write!(f, "SSH key Error: {e}"),
66+
ProtoError::UnsupportedCommand { command } => {
67+
write!(f, "Command not supported ({command})")
68+
}
6469
}
6570
}
6671
}

‎src/proto/message.rs

+21-14
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
use ssh_encoding::{CheckedSum, Decode, Encode, Error as EncodingError, Reader, Writer};
2-
use ssh_key::{private::KeypairData, public::KeyData, Error, Result, Signature};
2+
use ssh_key::{private::KeypairData, public::KeyData, Error, Signature};
3+
4+
use super::ProtoError;
5+
6+
type Result<T> = core::result::Result<T, ProtoError>;
37

48
#[derive(Clone, PartialEq, Debug)]
59
pub struct Identity {
@@ -21,7 +25,7 @@ impl Identity {
2125
}
2226

2327
impl Decode for Identity {
24-
type Error = Error;
28+
type Error = ProtoError;
2529

2630
fn decode(reader: &mut impl Reader) -> Result<Self> {
2731
let pubkey = reader.read_prefixed(KeyData::decode)?;
@@ -56,7 +60,7 @@ pub struct SignRequest {
5660
}
5761

5862
impl Decode for SignRequest {
59-
type Error = Error;
63+
type Error = ProtoError;
6064

6165
fn decode(reader: &mut impl Reader) -> Result<Self> {
6266
let pubkey = reader.read_prefixed(KeyData::decode)?;
@@ -97,7 +101,7 @@ pub struct AddIdentity {
97101
}
98102

99103
impl Decode for AddIdentity {
100-
type Error = Error;
104+
type Error = ProtoError;
101105

102106
fn decode(reader: &mut impl Reader) -> Result<Self> {
103107
let privkey = KeypairData::decode(reader)?;
@@ -126,7 +130,7 @@ pub struct AddIdentityConstrained {
126130
}
127131

128132
impl Decode for AddIdentityConstrained {
129-
type Error = Error;
133+
type Error = ProtoError;
130134

131135
fn decode(reader: &mut impl Reader) -> Result<Self> {
132136
let identity = AddIdentity::decode(reader)?;
@@ -168,7 +172,7 @@ pub struct RemoveIdentity {
168172
}
169173

170174
impl Decode for RemoveIdentity {
171-
type Error = Error;
175+
type Error = ProtoError;
172176

173177
fn decode(reader: &mut impl Reader) -> Result<Self> {
174178
let pubkey = reader.read_prefixed(KeyData::decode)?;
@@ -194,7 +198,7 @@ pub struct SmartcardKey {
194198
}
195199

196200
impl Decode for SmartcardKey {
197-
type Error = Error;
201+
type Error = ProtoError;
198202

199203
fn decode(reader: &mut impl Reader) -> Result<Self> {
200204
let id = String::decode(reader)?;
@@ -225,7 +229,7 @@ pub enum KeyConstraint {
225229
}
226230

227231
impl Decode for KeyConstraint {
228-
type Error = Error;
232+
type Error = ProtoError;
229233

230234
fn decode(reader: &mut impl Reader) -> Result<Self> {
231235
let constraint_type = u8::decode(reader)?;
@@ -239,7 +243,7 @@ impl Decode for KeyConstraint {
239243
reader.read(&mut details)?;
240244
KeyConstraint::Extension(name, details.into())
241245
}
242-
_ => return Err(Error::AlgorithmUnknown), // FIXME: it should be our own type
246+
_ => return Err(Error::AlgorithmUnknown)?, // FIXME: it should be our own type
243247
})
244248
}
245249
}
@@ -282,7 +286,7 @@ pub struct AddSmartcardKeyConstrained {
282286
}
283287

284288
impl Decode for AddSmartcardKeyConstrained {
285-
type Error = Error;
289+
type Error = ProtoError;
286290

287291
fn decode(reader: &mut impl Reader) -> Result<Self> {
288292
let key = SmartcardKey::decode(reader)?;
@@ -321,7 +325,7 @@ pub struct Extension {
321325
}
322326

323327
impl Decode for Extension {
324-
type Error = Error;
328+
type Error = ProtoError;
325329

326330
fn decode(reader: &mut impl Reader) -> Result<Self> {
327331
let name = String::decode(reader)?;
@@ -425,7 +429,7 @@ impl Message {
425429
}
426430

427431
impl Decode for Message {
428-
type Error = Error;
432+
type Error = ProtoError;
429433

430434
fn decode(reader: &mut impl Reader) -> Result<Self> {
431435
let message_type = u8::decode(reader)?;
@@ -436,7 +440,10 @@ impl Decode for Message {
436440
11 => Ok(Self::RequestIdentities),
437441
12 => Identity::decode_vec(reader).map(Self::IdentitiesAnswer),
438442
13 => SignRequest::decode(reader).map(Self::SignRequest),
439-
14 => reader.read_prefixed(|reader| Signature::decode(reader).map(Self::SignResponse)),
443+
14 => {
444+
Ok(reader
445+
.read_prefixed(|reader| Signature::decode(reader).map(Self::SignResponse))?)
446+
}
440447
17 => AddIdentity::decode(reader).map(Self::AddIdentity),
441448
18 => RemoveIdentity::decode(reader).map(Self::RemoveIdentity),
442449
19 => Ok(Self::RemoveAllIdentities),
@@ -448,7 +455,7 @@ impl Decode for Message {
448455
26 => AddSmartcardKeyConstrained::decode(reader).map(Self::AddSmartcardKeyConstrained),
449456
27 => Extension::decode(reader).map(Self::Extension),
450457
28 => Ok(Self::ExtensionFailure),
451-
_ => todo!(),
458+
command => Err(ProtoError::UnsupportedCommand { command }),
452459
}
453460
}
454461
}

0 commit comments

Comments
 (0)