1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168

pub mod icmp;
pub mod icmpv4_impl;
pub mod icmpv6_impl;
pub mod udp;
pub mod tcp;

use super::*;

use std::io;

///The possible headers on the transport layer
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum TransportHeader {
    Udp(udp::UdpHeader),
    Tcp(tcp::TcpHeader),
    Icmpv4(Icmpv4Header),
    Icmpv6(Icmpv6Header),
}

impl TransportHeader {

    /// Returns Result::Some containing the udp header if self has the value Udp.
    /// Otherwise None is returned.
    pub fn udp(self) -> Option<udp::UdpHeader> {
        use crate::TransportHeader::*;
        if let Udp(value) = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns Result::Some containing the udp header if self has the value Udp.
    /// Otherwise None is returned.
    pub fn mut_udp(&mut self) -> Option<&mut udp::UdpHeader> {
        use crate::TransportHeader::*;
        if let Udp(value) = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns Result::Some containing the tcp header if self has the value Tcp.
    /// Otherwise None is returned.
    pub fn tcp(self) -> Option<tcp::TcpHeader> {
        use crate::TransportHeader::*;
        if let Tcp(value)  = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns Result::Some containing a mutable refernce to the tcp header if self has the value Tcp.
    /// Otherwise None is returned.
    pub fn mut_tcp(&mut self) -> Option<&mut tcp::TcpHeader> {
        use crate::TransportHeader::*;
        if let Tcp(value) = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns Result::Some containing the ICMPv4 header if self has the value Icmpv4.
    /// Otherwise None is returned.
    pub fn icmpv4(self) -> Option<Icmpv4Header> {
        use crate::TransportHeader::*;
        if let Icmpv4(value) = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns Result::Some containing the ICMPv4 header if self has the value Icmpv4.
    /// Otherwise None is returned.
    pub fn mut_icmpv4(&mut self) -> Option<&mut Icmpv4Header> {
        use crate::TransportHeader::*;
        if let Icmpv4(value) = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns Result::Some containing the ICMPv6 header if self has the value Icmpv6.
    /// Otherwise None is returned.
    pub fn icmpv6(self) -> Option<Icmpv6Header> {
        use crate::TransportHeader::*;
        if let Icmpv6(value) = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns Result::Some containing the ICMPv6 header if self has the value Icmpv6.
    /// Otherwise None is returned.
    pub fn mut_icmpv6(&mut self) -> Option<&mut Icmpv6Header> {
        use crate::TransportHeader::*;
        if let Icmpv6(value) = self {
            Some(value)
        } else {
            None
        }
    }

    /// Returns the size of the transport header (in case of UDP fixed,
    /// in case of TCP cotanining the options).
    pub fn header_len(&self) -> usize {
        use crate::TransportHeader::*;
        match self {
            Udp(_) => udp::UdpHeader::SERIALIZED_SIZE,
            Tcp(value) => usize::from(value.header_len()),
            Icmpv4(value) => value.header_len(),
            Icmpv6(value) => value.header_len(),
        }
    }

    /// Calculates the checksum for the transport header & sets it in the header for
    /// an ipv4 header.
    pub fn update_checksum_ipv4(&mut self, ip_header: &Ipv4Header, payload: &[u8]) -> Result<(), ValueError> {
        use crate::TransportHeader::*;
        match self {
            Udp(header) => {
                header.checksum = header.calc_checksum_ipv4(ip_header, payload)?;
            },
            Tcp(header) => {
                header.checksum = header.calc_checksum_ipv4(ip_header, payload)?;
            },
            Icmpv4(header) => {
                header.update_checksum(payload);
            },
            Icmpv6(_) => return Err(ValueError::Icmpv6InIpv4),
        }
        Ok(())
    }

    /// Calculates the checksum for the transport header & sets it in the header for
    /// an ipv6 header.
    pub fn update_checksum_ipv6(&mut self, ip_header: &Ipv6Header, payload: &[u8]) -> Result<(), ValueError> {
        use crate::TransportHeader::*;
        match self {
            Icmpv4(header) => header.update_checksum(payload),
            Icmpv6(header) => header.update_checksum(ip_header.source, ip_header.destination, payload)?,
            Udp(header) => {
                header.checksum = header.calc_checksum_ipv6(ip_header, payload)?;
            },
            Tcp(header) => {
                header.checksum = header.calc_checksum_ipv6(ip_header, payload)?;
            }
        }
        Ok(())
    }

    /// Write the transport header to the given writer.
    pub fn write<T: io::Write + Sized>(&self, writer: &mut T) -> Result<(), WriteError> {
        use crate::TransportHeader::*;
        match self {
            Icmpv4(value) => value.write(writer),
            Icmpv6(value) => value.write(writer),
            Udp(value) => value.write(writer),
            Tcp(value) => value.write(writer).map_err(WriteError::from)
        }
    }
}