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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183

use super::super::*;

extern crate byteorder;
use self::byteorder::{ByteOrder, BigEndian, ReadBytesExt, WriteBytesExt};

///Udp header according to rfc768.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct UdpHeader {
    ///Source port of the packet (optional).
    pub source_port: u16,
    ///Destination port of the packet.
    pub destination_port: u16,
    ///Length of the packet (includes the udp header length of 8 bytes).
    pub length: u16,
    ///The checksum of the packet. The checksum is calculated from a pseudo header, the udp header and the payload. The pseudo header is composed of source and destination address, protocol number 
    pub checksum: u16
}

impl UdpHeader {

    ///Returns an udp header for the given parameters
    pub fn without_ipv4_checksum(source_port: u16, destination_port: u16, payload_length: usize) -> Result<UdpHeader, ValueError> {
        //check that the total length fits into the field
        const MAX_PAYLOAD_LENGTH: usize = (std::u16::MAX as usize) - UdpHeader::SERIALIZED_SIZE;
        if MAX_PAYLOAD_LENGTH < payload_length {
            return Err(ValueError::UdpPayloadLengthTooLarge(payload_length));
        }

        Ok(UdpHeader{
            source_port: source_port,
            destination_port: destination_port,
            length: (UdpHeader::SERIALIZED_SIZE + payload_length) as u16, //payload plus udp header
            checksum: 0
        })
    }

    ///Calculate an udp header given an ipv4 header and the payload
    pub fn with_ipv4_checksum(source_port: u16, destination_port: u16, ip_header: &Ipv4Header, payload: &[u8]) -> Result<UdpHeader, ValueError> {

        //check that the total length fits into the field
        const MAX_PAYLOAD_LENGTH: usize = (std::u16::MAX as usize) - UdpHeader::SERIALIZED_SIZE;
        if MAX_PAYLOAD_LENGTH < payload.len() {
            return Err(ValueError::UdpPayloadLengthTooLarge(payload.len()));
        }

        let mut result = UdpHeader{
            source_port: source_port,
            destination_port: destination_port,
            length: (UdpHeader::SERIALIZED_SIZE + payload.len()) as u16, //payload plus udp header
            checksum: 0
        };
        result.checksum = result.calc_checksum_ipv4_internal(&ip_header.source, &ip_header.destination, ip_header.protocol, payload);
        Ok(result)
    }

    ///Calculates the upd header checksum based on a ipv4 header.
    pub fn calc_checksum_ipv4(&self, ip_header: &Ipv4Header, payload: &[u8]) -> Result<u16, ValueError> {
        self.calc_checksum_ipv4_raw(&ip_header.source, &ip_header.destination, ip_header.protocol, payload)
    }

    ///Calculates the upd header checksum based on a ipv4 header.
    pub fn calc_checksum_ipv4_raw(&self, source: &[u8;4], destination: &[u8;4], protocol: u8, payload: &[u8]) -> Result<u16, ValueError> {
        //check that the total length fits into the field
        const MAX_PAYLOAD_LENGTH: usize = (std::u16::MAX as usize) - UdpHeader::SERIALIZED_SIZE;
        if MAX_PAYLOAD_LENGTH < payload.len() {
            return Err(ValueError::UdpPayloadLengthTooLarge(payload.len()));
        }

        Ok(self.calc_checksum_ipv4_internal(source, destination, protocol, payload))
    }
    
    ///Calculates the upd header checksum based on a ipv4 header.
    fn calc_checksum_ipv4_internal(&self, source: &[u8;4], destination: &[u8;4], protocol: u8, payload: &[u8]) -> u16 {
        self.calc_checksum_post_ip(BigEndian::read_u16(&source[0..2]) as u64 + //pseudo header
                                   BigEndian::read_u16(&source[2..4]) as u64 +
                                   BigEndian::read_u16(&destination[0..2]) as u64 +
                                   BigEndian::read_u16(&destination[2..4]) as u64 +
                                   protocol as u64 +
                                   self.length as u64, 
                                   payload)
    }

    ///Calculate an udp header given an ipv6 header and the payload
    pub fn with_ipv6_checksum(source_port: u16, destination_port: u16, ip_header: &Ipv6Header, payload: &[u8]) -> Result<UdpHeader, ValueError> {

        //check that the total length fits into the field
        const MAX_PAYLOAD_LENGTH: usize = (std::u16::MAX as usize) - UdpHeader::SERIALIZED_SIZE;
        if MAX_PAYLOAD_LENGTH < payload.len() {
            return Err(ValueError::UdpPayloadLengthTooLarge(payload.len()));
        }

        let mut result = UdpHeader{
            source_port: source_port,
            destination_port: destination_port,
            length: (UdpHeader::SERIALIZED_SIZE + payload.len()) as u16, //payload plus udp header
            checksum: 0
        };
        result.checksum = result.calc_checksum_ipv6_internal(&ip_header.source, &ip_header.destination, payload);
        Ok(result)
    }

    ///Calculates the checksum of the current udp header given an ipv6 header and the payload.
    pub fn calc_checksum_ipv6(&self, ip_header: &Ipv6Header, payload: &[u8]) -> Result<u16, ValueError> {
        self.calc_checksum_ipv6_raw(&ip_header.source, &ip_header.destination, payload)
    }

    ///Calculates the checksum of the current udp header given an ipv6 source & destination address plus the payload.
    pub fn calc_checksum_ipv6_raw(&self, source: &[u8;16], destination: &[u8;16], payload: &[u8]) -> Result<u16, ValueError> {
        //check that the total length fits into the field
        const MAX_PAYLOAD_LENGTH: usize = (std::u16::MAX as usize) - UdpHeader::SERIALIZED_SIZE;
        if MAX_PAYLOAD_LENGTH < payload.len() {
            return Err(ValueError::UdpPayloadLengthTooLarge(payload.len()));
        }

        Ok(self.calc_checksum_ipv6_internal(source, destination, payload))
    }

    fn calc_checksum_ipv6_internal(&self, source: &[u8;16], destination: &[u8;16], payload: &[u8]) -> u16 {
        fn calc_sum(value: &[u8;16]) -> u64 {
            let mut result = 0;
            for i in 0..8 {
                let index = i*2;
                result += BigEndian::read_u16(&value[index..(index + 2)]) as u64;
            }
            result
        }
        self.calc_checksum_post_ip(calc_sum(source) +
                                   calc_sum(destination) +
                                   IpTrafficClass::Udp as u64 +
                                   self.length as u64,
                                   payload)
    }

    ///This method takes the sum of the preudo ip header and calculates the rest of the checksum.
    fn calc_checksum_post_ip(&self, ip_pseudo_header_sum: u64, payload: &[u8]) -> u16 {
        let mut sum = ip_pseudo_header_sum +
                      self.source_port as u64 + //udp header start
                      self.destination_port as u64 +
                      self.length as u64;

        for i in 0..(payload.len()/2) {
            sum += BigEndian::read_u16(&payload[i*2..i*2 + 2]) as u64;
        }
        //pad the last byte with 0
        if payload.len() % 2 == 1 {
            sum += BigEndian::read_u16(&[*payload.last().unwrap(), 0]) as u64;
        }
        let carry_add = (sum & 0xffff) + 
                        ((sum >> 16) & 0xffff) +
                        ((sum >> 32) & 0xffff) +
                        ((sum >> 48) & 0xffff);
        let result = ((carry_add & 0xffff) + (carry_add >> 16)) as u16;
        if 0xffff == result {
            result //avoid the transmition of an all 0 checksum as this value is reserved by "checksum disabled" (see rfc)
        } else {
            !result
        }
    }

    ///Tries to read an udp header from the current position.
    pub fn read<T: io::Read + io::Seek + Sized>(reader: &mut T) -> Result<UdpHeader, io::Error> {
        Ok(UdpHeader{
            source_port: reader.read_u16::<BigEndian>()?,
            destination_port: reader.read_u16::<BigEndian>()?,
            length: reader.read_u16::<BigEndian>()?,
            checksum: reader.read_u16::<BigEndian>()?
        })
    }

    ///Write the udp header without recalculating the checksum or length.
    pub fn write<T: io::Write + Sized>(&self, writer: &mut T) -> Result<(), WriteError> {
        writer.write_u16::<BigEndian>(self.source_port)?;
        writer.write_u16::<BigEndian>(self.destination_port)?;
        writer.write_u16::<BigEndian>(self.length)?;
        writer.write_u16::<BigEndian>(self.checksum)?;
        Ok(())
    }
}

impl SerializedSize for UdpHeader {
    ///Size of the header itself
    const SERIALIZED_SIZE: usize = 8;
}