001    /*
002     * Licensed to the Apache Software Foundation (ASF) under one or more
003     * contributor license agreements.  See the NOTICE file distributed with
004     * this work for additional information regarding copyright ownership.
005     * The ASF licenses this file to You under the Apache License, Version 2.0
006     * (the "License"); you may not use this file except in compliance with
007     * the License.  You may obtain a copy of the License at
008     *
009     *      http://www.apache.org/licenses/LICENSE-2.0
010     *
011     * Unless required by applicable law or agreed to in writing, software
012     * distributed under the License is distributed on an "AS IS" BASIS,
013     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014     * See the License for the specific language governing permissions and
015     * limitations under the License.
016     */
017    package org.apache.commons.net.tftp;
018    
019    import java.io.BufferedInputStream;
020    import java.io.BufferedOutputStream;
021    import java.io.File;
022    import java.io.FileInputStream;
023    import java.io.FileOutputStream;
024    import java.io.IOException;
025    import java.io.InputStream;
026    import java.io.OutputStream;
027    
028    import org.apache.commons.net.tftp.TFTPServer.ServerMode;
029    
030    import junit.framework.TestCase;
031    
032    /**
033     * Test the TFTP Server and TFTP Client by creating some files in the system temp folder and then
034     * uploading and downloading them.
035     * 
036     * @author <A HREF="mailto:daniel.armbrust.list@gmail.com">Dan Armbrust</A>
037     * 
038     */
039    public class TFTPTest extends TestCase
040    {
041        static TFTPServer tftpS;
042        static File serverDirectory = new File(System.getProperty("java.io.tmpdir"));
043        static String filePrefix = "tftp-";
044        static File[] files = new File[8];
045    
046        static int testsLeftToRun = 6;
047    
048        // only want to do this once...
049        static
050        {
051            try
052            {
053                files[0] = createFile(new File(serverDirectory, filePrefix + "empty.txt"), 0);
054                files[1] = createFile(new File(serverDirectory, filePrefix + "small.txt"), 1);
055                files[2] = createFile(new File(serverDirectory, filePrefix + "511.txt"), 511);
056                files[3] = createFile(new File(serverDirectory, filePrefix + "512.txt"), 512);
057                files[4] = createFile(new File(serverDirectory, filePrefix + "513.txt"), 513);
058                files[5] = createFile(new File(serverDirectory, filePrefix + "med.txt"), 1000 * 1024);
059                files[6] = createFile(new File(serverDirectory, filePrefix + "big.txt"), 5000 * 1024);
060                files[7] = createFile(new File(serverDirectory, filePrefix + "huge.txt"), 37000 * 1024);
061    
062                // Start the server
063                tftpS = new TFTPServer(serverDirectory, serverDirectory, 6900, ServerMode.GET_AND_PUT,
064                        null, null);
065                tftpS.setSocketTimeout(2000);
066            }
067            catch (IOException e)
068            {
069                e.printStackTrace();
070            }
071    
072        }
073    
074        @Override
075        protected void tearDown() throws Exception
076        {
077            testsLeftToRun--;
078            if (testsLeftToRun <= 0)
079            {
080                if (tftpS != null)
081                {
082                    tftpS.shutdown();
083                }
084                for (int i = 0; i < files.length; i++)
085                {
086                    files[i].delete();
087                }
088            }
089            super.tearDown();
090        }
091    
092        /*
093         * Create a file, size specified in bytes
094         */
095        private static File createFile(File file, int size) throws IOException
096        {
097            OutputStream os = new BufferedOutputStream(new FileOutputStream(file));
098            byte[] temp = "0".getBytes();
099            for (int i = 0; i < size; i++)
100            {
101                os.write(temp);
102            }
103            os.close();
104            return file;
105        }
106    
107        public void testTFTPBinaryDownloads() throws Exception
108        {
109            // test with the smaller files
110            for (int i = 0; i < 6; i++)
111            {
112                testDownload(TFTP.BINARY_MODE, files[i]);
113            }
114        }
115    
116        public void testASCIIDownloads() throws Exception
117        {
118            // test with the smaller files
119            for (int i = 0; i < 6; i++)
120            {
121                testDownload(TFTP.ASCII_MODE, files[i]);
122            }
123        }
124    
125        public void testTFTPBinaryUploads() throws Exception
126        {
127            // test with the smaller files
128            for (int i = 0; i < 6; i++)
129            {
130                testUpload(TFTP.BINARY_MODE, files[i]);
131            }
132        }
133    
134        public void testASCIIUploads() throws Exception
135        {
136            // test with the smaller files
137            for (int i = 0; i < 6; i++)
138            {
139                testUpload(TFTP.ASCII_MODE, files[i]);
140            }
141        }
142    
143        public void testHugeUploads() throws Exception
144        {
145            for (int i = 5; i < files.length; i++)
146            {
147                testUpload(TFTP.BINARY_MODE, files[i]);
148            }
149        }
150    
151        public void testHugeDownloads() throws Exception
152        {
153            // test with the smaller files
154            for (int i = 5; i < files.length; i++)
155            {
156                testDownload(TFTP.BINARY_MODE, files[i]);
157            }
158        }
159    
160        private void testDownload(int mode, File file) throws IOException
161        {
162            // Create our TFTP instance to handle the file transfer.
163            TFTPClient tftp = new TFTPClient();
164            tftp.open();
165            tftp.setSoTimeout(2000);
166    
167            File out = new File(serverDirectory, filePrefix + "download");
168    
169            // cleanup old failed runs
170            out.delete();
171            assertTrue("Couldn't clear output location", !out.exists());
172    
173            FileOutputStream output = new FileOutputStream(out);
174    
175            tftp.receiveFile(file.getName(), mode, output, "localhost", 6900);
176            output.close();
177    
178            assertTrue("file not created", out.exists());
179            assertTrue("files not identical on file " + file, filesIdentical(out, file));
180    
181            // delete the downloaded file
182            out.delete();
183        }
184    
185        private void testUpload(int mode, File file) throws Exception
186        {
187            // Create our TFTP instance to handle the file transfer.
188            TFTPClient tftp = new TFTPClient();
189            tftp.open();
190            tftp.setSoTimeout(2000);
191    
192            File in = new File(serverDirectory, filePrefix + "upload");
193            // cleanup old failed runs
194            in.delete();
195            assertTrue("Couldn't clear output location", !in.exists());
196    
197            FileInputStream fis = new FileInputStream(file);
198            tftp.sendFile(in.getName(), mode, fis, "localhost", 6900);
199            fis.close();
200    
201            // need to give the server a bit of time to receive our last packet, and
202            // close out its file buffers, etc.
203            Thread.sleep(100);
204            assertTrue("file not created", in.exists());
205            assertTrue("files not identical on file " + file, filesIdentical(file, in));
206    
207            in.delete();
208        }
209    
210        private boolean filesIdentical(File a, File b) throws IOException
211        {
212            if (!a.exists() || !b.exists())
213            {
214                return false;
215            }
216    
217            if (a.length() != b.length())
218            {
219                return false;
220            }
221    
222            InputStream fisA = new BufferedInputStream(new FileInputStream(a));
223            InputStream fisB = new BufferedInputStream(new FileInputStream(b));
224    
225            int aBit = fisA.read();
226            int bBit = fisB.read();
227    
228            while (aBit != -1)
229            {
230                if (aBit != bBit)
231                {
232                    fisA.close();
233                    fisB.close();
234                    return false;
235                }
236                aBit = fisA.read();
237                bBit = fisB.read();
238            }
239    
240            fisA.close();
241            fisB.close();
242            return true;
243        }
244    }