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
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
use async_io::Async;
use async_lock::Lock;
use bytes::BytesMut;
use std::collections::HashSet;
use std::net::{SocketAddr, UdpSocket};
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;

use super::handlers::{DirHandler, DirHandlerMode};
use super::{Handler, ServerConfig, TftpServer};
use crate::error::{Error, Result};

/// TFTP server builder.
pub struct TftpServerBuilder<H: Handler> {
    handle: H,
    addr: SocketAddr,
    socket: Option<Async<UdpSocket>>,
    timeout: Duration,
    block_size_limit: Option<u16>,
    max_send_retries: u32,
    ignore_client_timeout: bool,
    ignore_client_block_size: bool,
}

impl TftpServerBuilder<DirHandler> {
    /// Create new buidler with [`DirHandler`] that serves only read requests.
    ///
    /// [`DirHandler`]: handlers/struct.DirHandler.html
    pub fn with_dir_ro<P>(dir: P) -> Result<TftpServerBuilder<DirHandler>>
    where
        P: AsRef<Path>,
    {
        let handler = DirHandler::new(dir, DirHandlerMode::ReadOnly)?;
        Ok(TftpServerBuilder::with_handler(handler))
    }

    /// Create new buidler with [`DirHandler`] that serves only write requests.
    ///
    /// [`DirHandler`]: handlers/struct.DirHandler.html
    pub fn with_dir_wo<P>(dir: P) -> Result<TftpServerBuilder<DirHandler>>
    where
        P: AsRef<Path>,
    {
        let handler = DirHandler::new(dir, DirHandlerMode::WriteOnly)?;
        Ok(TftpServerBuilder::with_handler(handler))
    }

    /// Create new buidler with [`DirHandler`] that serves read and write requests.
    ///
    /// [`DirHandler`]: handlers/struct.DirHandler.html
    pub fn with_dir_rw<P>(dir: P) -> Result<TftpServerBuilder<DirHandler>>
    where
        P: AsRef<Path>,
    {
        let handler = DirHandler::new(dir, DirHandlerMode::ReadWrite)?;
        Ok(TftpServerBuilder::with_handler(handler))
    }
}

impl<H: Handler> TftpServerBuilder<H> {
    /// Create new builder with custom [`Handler`].
    ///
    /// [`Handler`]: trait.Handler.html
    pub fn with_handler(handler: H) -> Self {
        TftpServerBuilder {
            handle: handler,
            addr: "0.0.0.0:69".parse().unwrap(),
            socket: None,
            timeout: Duration::from_secs(3),
            block_size_limit: None,
            max_send_retries: 100,
            ignore_client_timeout: false,
            ignore_client_block_size: false,
        }
    }

    /// Set listening address.
    ///
    /// This is ignored if underling socket is set.
    ///
    /// **Default:** `0.0.0.0:69`
    pub fn bind(self, addr: SocketAddr) -> Self {
        TftpServerBuilder {
            addr,
            ..self
        }
    }

    /// Set underling UDP socket.
    pub fn socket(self, socket: Async<UdpSocket>) -> Self {
        TftpServerBuilder {
            socket: Some(socket),
            ..self
        }
    }

    /// Set underling UDP socket.
    pub fn std_socket(self, socket: UdpSocket) -> Result<Self> {
        let socket = Async::new(socket)?;

        Ok(TftpServerBuilder {
            socket: Some(socket),
            ..self
        })
    }

    /// Set retry timeout.
    ///
    /// Client can override this (RFC2349). If you want to enforce it you must
    /// combine it [`ignore_client_timeout`].
    ///
    /// This crate allows you to set non-standard timeouts (i.e. timeouts that are less
    /// than a second). However if you choose to do it make sure you test it well in your
    /// environment since client's behavior is undefined.
    ///
    /// **Default:** 3 seconds
    ///
    /// [`ignore_client_timeout`]: struct.TftpServerBuilder.html#method.ignore_client_timeout
    pub fn timeout(self, timeout: Duration) -> Self {
        TftpServerBuilder {
            timeout,
            ..self
        }
    }

    /// Set maximum block size.
    ///
    /// Client can request a specific block size (RFC2348). Use this option if you
    /// want to set a limit.
    ///
    /// **Real life scenario:** U-Boot does not support IP fragmentation and requests
    /// block size of 1468. This works fine if your MTU is 1500 bytes, however if
    /// you are accessing client through a VPN, then transfer will never start. Use
    /// this option to workaround the problem.
    pub fn block_size_limit(self, size: u16) -> Self {
        TftpServerBuilder {
            block_size_limit: Some(size),
            ..self
        }
    }

    /// Set maximum send retries for a data block.
    ///
    /// On timeout server will try to send the data block again. When retries are
    /// reached for the specific data block the server closes the connection with
    /// the client.
    ///
    /// Default: 100 retries.
    pub fn max_send_retries(self, retries: u32) -> Self {
        TftpServerBuilder {
            max_send_retries: retries,
            ..self
        }
    }

    /// Ignore client's `timeout` option.
    ///
    /// With this you enforce server's timeout by ignoring client's
    /// `timeout` option of RFC2349.
    pub fn ignore_client_timeout(self) -> Self {
        TftpServerBuilder {
            ignore_client_timeout: true,
            ..self
        }
    }

    /// Ignore client's block size option.
    ///
    /// With this you can ignore client's `blksize` option of RFC2348.
    /// This will enforce 512 block size that is defined in RFC1350.
    pub fn ignore_client_block_size(self) -> Self {
        TftpServerBuilder {
            ignore_client_block_size: true,
            ..self
        }
    }

    /// Build [`TftpServer`].
    ///
    /// [`TftpServer`]: struct.TftpServer.html
    pub async fn build(mut self) -> Result<TftpServer<H>> {
        let socket = match self.socket.take() {
            Some(socket) => socket,
            None => Async::<UdpSocket>::bind(self.addr).map_err(Error::Bind)?,
        };

        let config = ServerConfig {
            timeout: self.timeout,
            block_size_limit: self.block_size_limit,
            max_send_retries: self.max_send_retries,
            ignore_client_timeout: self.ignore_client_timeout,
            ignore_client_block_size: self.ignore_client_block_size,
        };

        Ok(TftpServer {
            socket: Some(socket),
            handler: Arc::new(Lock::new(self.handle)),
            config,
            reqs_in_progress: HashSet::new(),
            buffer: BytesMut::new(),
        })
    }
}