2 files changed, 97 insertions, 59 deletions

diff --git a/modules/es_ftw/es_ftw.cpp b/modules/es_ftw/es_ftw.cpp
index fd5ec3fde..d13016104 100644
--- a/modules/es_ftw/es_ftw.cpp
+++ b/modules/es_ftw/es_ftw.cpp
@@ -6,7 +6,7 @@
 #include <botan/es_ftw.h>
 #include <botan/util.h>
 #include <cstring>
-#include <vector>
+#include <deque>
 
 #ifndef _POSIX_C_SOURCE
   #define _POSIX_C_SOURCE 199309
@@ -20,6 +20,79 @@
 
 namespace Botan {
 
+namespace {
+
+class Directory_Walker
+   {
+   public:
+      Directory_Walker(const std::string& root) { add_directory(root); }
+      ~Directory_Walker();
+
+      int next_fd();
+   private:
+      void add_directory(const std::string&);
+
+      std::deque<std::pair<DIR*, std::string> > dirs;
+   };
+
+void Directory_Walker::add_directory(const std::string& dirname)
+   {
+   DIR* dir = ::opendir(dirname.c_str());
+   if(dir)
+      dirs.push_back(std::make_pair(dir, dirname));
+   }
+
+Directory_Walker::~Directory_Walker()
+   {
+   while(dirs.size())
+      {
+      ::closedir(dirs[0].first);
+      dirs.pop_front();
+      }
+   }
+
+int Directory_Walker::next_fd()
+   {
+   while(dirs.size())
+      {
+      std::pair<DIR*, std::string> dirinfo = dirs[0];
+
+      struct dirent* entry = ::readdir(dirinfo.first);
+
+      if(!entry)
+         {
+         ::closedir(dirinfo.first);
+         dirs.pop_front();
+         continue;
+         }
+
+      const std::string filename = entry->d_name;
+
+      if(filename == "." || filename == "..")
+         continue;
+
+      const std::string full_path = dirinfo.second + '/' + filename;
+
+      struct stat stat_buf;
+      if(::lstat(full_path.c_str(), &stat_buf) == -1)
+         continue;
+
+      if(S_ISDIR(stat_buf.st_mode))
+         add_directory(full_path);
+      else if(S_ISREG(stat_buf.st_mode))
+         {
+         int fd = ::open(full_path.c_str(), O_RDONLY | O_NOCTTY);
+
+         if(fd > 0)
+            return fd;
+         }
+      }
+
+   return -1;
+   }
+
+}
+
 /*************************************************
 * FTW_EntropySource Constructor                  *
 *************************************************/
@@ -32,9 +105,7 @@ FTW_EntropySource::FTW_EntropySource(const std::string& p) : path(p)
 *************************************************/
 void FTW_EntropySource::do_fast_poll()
    {
-   files_read = 0;
-   max_read = 32;
-   gather_from_dir(path);
+   poll(32*1024);
    }
 
 /*************************************************
@@ -42,67 +113,34 @@ void FTW_EntropySource::do_fast_poll()
 *************************************************/
 void FTW_EntropySource::do_slow_poll()
    {
-   files_read = 0;
-   max_read = 256;
-   gather_from_dir(path);
+   poll(256*1024);
    }
 
 /*************************************************
-* Gather Entropy From Directory Tree             *
+* FTW Poll                                       *
 *************************************************/
-void FTW_EntropySource::gather_from_dir(const std::string& dirname)
+void FTW_EntropySource::poll(u32bit max_read)
    {
-   if(dirname == "" || files_read >= max_read)
-      return;
-
-   DIR* dir = ::opendir(dirname.c_str());
-   if(dir == 0)
-      return;
+   Directory_Walker dir(path);
+   u32bit read_so_far = 0;
 
-   std::vector<std::string> subdirs;
-
-   dirent* entry = ::readdir(dir);
-   while(entry && (files_read < max_read))
+   while(read_so_far < max_read)
       {
-      if((std::strcmp(entry->d_name, ".") == 0) ||
-         (std::strcmp(entry->d_name, "..") == 0))
-         { entry = ::readdir(dir); continue; }
-
-      const std::string filename = dirname + '/' + entry->d_name;
-
-      struct stat stat_buf;
-      if(::lstat(filename.c_str(), &stat_buf) == -1)
-         { entry = ::readdir(dir); continue; }
-
-      if(S_ISREG(stat_buf.st_mode))
-         gather_from_file(filename);
-      else if(S_ISDIR(stat_buf.st_mode))
-         subdirs.push_back(filename);
-      entry = ::readdir(dir);
-      }
-   ::closedir(dir);
+      int fd = dir.next_fd();
 
-   for(u32bit j = 0; j != subdirs.size(); j++)
-      gather_from_dir(subdirs[j]);
-   }
+      if(fd == -1)
+         break;
 
-/*************************************************
-* Gather Entropy From A File                     *
-*************************************************/
-void FTW_EntropySource::gather_from_file(const std::string& filename)
-   {
-   int fd = ::open(filename.c_str(), O_RDONLY | O_NOCTTY);
-   if(fd == -1)
-      return;
+      SecureVector<byte> read_buf(1024);
+      ssize_t got = ::read(fd, read_buf.begin(), read_buf.size());
 
-   SecureVector<byte> read_buf(1024);
-   ssize_t got = ::read(fd, read_buf.begin(), read_buf.size());
-   ::close(fd);
+      if(got > 0)
+         {
+         add_bytes(read_buf, got);
+         read_so_far += got;
+         }
 
-   if(got > 0)
-      {
-      add_bytes(read_buf, got);
-      files_read++;
+      ::close(fd);
       }
    }
 
diff --git a/modules/es_ftw/es_ftw.h b/modules/es_ftw/es_ftw.h
index 15d113890..0a48cc551 100644
--- a/modules/es_ftw/es_ftw.h
+++ b/modules/es_ftw/es_ftw.h
@@ -1,6 +1,6 @@
 /*************************************************
 * File Tree Walking EntropySource Header File    *
-* (C) 1999-2007 Jack Lloyd                       *
+* (C) 1999-2008 Jack Lloyd                       *
 *************************************************/
 
 #ifndef BOTAN_EXT_ENTROPY_SRC_FTW_H__
@@ -16,14 +16,14 @@ namespace Botan {
 class FTW_EntropySource : public Buffered_EntropySource
    {
    public:
-      FTW_EntropySource(const std::string& = "/proc");
+      FTW_EntropySource(const std::string& root_dir);
    private:
       void do_fast_poll();
       void do_slow_poll();
-      void gather_from_dir(const std::string&);
-      void gather_from_file(const std::string&);
+
+      void poll(u32bit max_read);
+
       const std::string path;
-      u32bit files_read, max_read;
    };
 
 }