sftp_readahead.c

/*
 * Example that demonstrates SFTP read-ahead with libssh
 *
 * Usage:
 * sftp_readahead <user> <password> <host> <filepath> [<read_ahead_size>]
 */

#include <ctype.h>
#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <libssh/libssh.h>
#include <libssh/sftp.h>
#include <sys/stat.h>
#include <fcntl.h>

// Declarations

typedef struct {
   int* buffer;
   int* buffer_end;
   int* data_start;
   int* data_end;
   uint64_t  count;
   uint64_t  size;
} ring_buffer;

static int sftp_fread(char* buffer, size_t bytesRequested, sftp_file handle, ring_buffer* readAheadBuf, size_t* bytesReturned);

static int  ringbuffer_init(ring_buffer* rb, uint64_t size);
static void ringbuffer_free(ring_buffer* rb);
static void ringbuffer_push(ring_buffer* rb, int data);
static int  ringbuffer_pop(ring_buffer* rb);
static int  ringbuffer_full(ring_buffer* rb);
static int  ringbuffer_empty(ring_buffer* rb);

// Implementation

static int sftp_fread(char* buffer, size_t bytesRequested, sftp_file handle, ring_buffer* readAheadBuf, size_t* bytesReturned) {
   size_t bytesToRead = bytesRequested;
   int bytesRead = 1, async_request, result = 0;

   // ensure fread() semantics (always return full buffers unless eof/error)
   while (bytesToRead > 0) {
      if (bytesRead > 0) {
         // send a request and put its handle into the ring buffer until full
         while (!ringbuffer_full(readAheadBuf)) {
            async_request = sftp_async_read_begin(handle, bytesToRead);
            //fprintf(stderr, "Sent request %d\n", async_request);
            if (async_request >= 0) {
               ringbuffer_push(readAheadBuf, async_request);
            }
         }
      }
      // read data from one response
      if (!ringbuffer_empty(readAheadBuf)) {
         async_request = ringbuffer_pop(readAheadBuf);
         //fprintf(stderr, "Waiting for reply %d\n", async_request);
         bytesRead = sftp_async_read(handle, buffer, bytesToRead, async_request);
         if (bytesRead>0) {
            bytesToRead -= bytesRead;
            buffer+=bytesRead;
         } else if (bytesRead == 0) { // EOF
            result = 1;
            break;
         } else {  // Error
            result = 2;
            break;
         }
      } else {
         fprintf(stderr, "Error: ringbuffer empty\n");
         result = 2;
         break;
      }
   }

   *bytesReturned = bytesRequested - bytesToRead;
   return result;
}

int main(int argc, char *argv[]) {
   char *username=NULL, *password=NULL, *hostname=NULL, *filepath=NULL;
   ssh_session sshSession;
   sftp_session sftpSession;
   sftp_file sftpFile;
   ring_buffer readAheadBuf;
   long timeout = 30;
   int rc;
   char buffer[4096];
   int readAheadMax = 128;
   size_t read, readTotal = 0;

   if (argc != 5 && argc != 6) {
      fprintf(stderr, "Usage: %s <user> <password> <host> <filepath> [<read_ahead_size>]\n", argv[0]);
      fprintf(stderr, "<read_ahead_size> defaults to %d blocks of 4 KiB\n", readAheadMax);
      return 1;
   }
   username = argv[1];
   password = argv[2];
   hostname = argv[3];
   filepath = argv[4];
   if (argc == 6) {
      readAheadMax = atoi(argv[5]);
   }

   if (readAheadMax <= 0) {
      fprintf(stderr, "<read_ahead_size> parameter must be a positive number of 4 KiB blocks\n");
   }

   // create SSH session handle
   sshSession = ssh_new();
   if (sshSession == NULL) {
      fprintf(stderr, "libssh failed to initialize the session!\n");
      return 1;
   }

   // configure session
   ssh_options_set(sshSession, SSH_OPTIONS_HOST, hostname);
   ssh_options_set(sshSession, SSH_OPTIONS_USER, username);
   ssh_options_set(sshSession, SSH_OPTIONS_TIMEOUT, &timeout);
   ssh_options_set(sshSession, SSH_OPTIONS_LOG_VERBOSITY, "SSH_LOG_PACKET");

   // Connect to server
   rc = ssh_connect(sshSession);
   if (rc != SSH_OK) {
      fprintf(stderr, "Error connecting to %s: %s\n", hostname, ssh_get_error(sshSession));
      return 1;
   }

   fprintf(stderr, "Connected.\n");

   // set blocking mode
   ssh_set_blocking(sshSession, 1);

   // authenticate with password
   rc = ssh_userauth_password(sshSession, NULL, password);
   if (rc != SSH_AUTH_SUCCESS) {
      fprintf(stderr, "Error authenticating with password: %s\n", ssh_get_error(sshSession));
      return 1;
   }

   fprintf(stderr, "Authenticated.\n");

   // start SFTP session
   sftpSession = sftp_new(sshSession);
   if (sftpSession == NULL) {
       fprintf(stderr, "Unable to init SFTP session\n");
       return 1;
   }

   rc = sftp_init(sftpSession);
   if (rc != SSH_OK) {
      fprintf(stderr, "Error initializing SFTP session: %s.\n", ssh_get_error(sshSession));
      return 1;
   }

   // open file via SFTP
   sftpFile =  sftp_open(sftpSession, filepath, O_RDONLY, 0);
   if (sftpFile == NULL) {
      fprintf(stderr, "Cannot open file %s for reading: %s\n", filepath, ssh_get_error(sshSession));
      return 1;
   }

   // initialize ring buffer
   if (!ringbuffer_init(&readAheadBuf, readAheadMax)) {
      fprintf(stderr, "Could not allocate read ahead buffer\n");
      return 1;
   }

   fprintf(stderr, "Reading file %s ...\n", filepath);

   // read the whole file
   do {
      rc = sftp_fread(buffer, sizeof(buffer), sftpFile, &readAheadBuf, &read);
      readTotal += read;
   } while (rc == 0);

   if (rc == 1) {
      fprintf(stderr, "EOF reached\n");
   } else if (rc == 2) {
      fprintf(stderr, "Error while reading: %s\n", ssh_get_error(sshSession));
   }

   fprintf(stderr, "Read %lu bytes.\n", (unsigned long)readTotal);

   /* Need to read any requested blocks prior to sftp_close()?
   while (!ringbuffer_empty(&readAheadBuf)) {
      // skip over any requested but unfetched data
      rc = ringbuffer_pop(&readAheadBuf);
      fprintf(stderr, "skipping %d\n", rc);
      sftp_async_read(sftpFile, buffer, bufsize, rc);
   }
   */

   // clean up
   fprintf(stderr, "sftp_close()");
   rc = sftp_close(sftpFile);
   fprintf(stderr, " done\n");
   if(rc==SSH_ERROR) {
      fprintf(stderr, "Failed to close the file\n");
   }
   sftp_free(sftpSession);
   ssh_disconnect(sshSession);
   ssh_free(sshSession);
   ringbuffer_free(&readAheadBuf);

   fprintf(stderr, "Done\n");

   return 0;
}

// *********************Ring Buffer implementation**********************
static int ringbuffer_init(ring_buffer* rb, uint64_t size) {
   rb->buffer = malloc(sizeof(int) * size);
   rb->buffer_end = rb->buffer + size;
   rb->size = size;
   rb->data_start = rb->buffer;
   rb->data_end = rb->buffer;
   rb->count = 0;
   return rb->buffer != NULL;
}

static void ringbuffer_free(ring_buffer* rb) {
   if (rb->buffer != NULL) {
      free(rb->buffer);
      rb->buffer = NULL;
   }
}

static void ringbuffer_push(ring_buffer* rb, int data) {
   *rb->data_end = data;
   rb->data_end++;
   if (rb->data_end == rb->buffer_end)
      rb->data_end = rb->buffer;

   if (ringbuffer_full(rb)) {
      if ((rb->data_start + 1) == rb->buffer_end)
         rb->data_start = rb->buffer;
      else
         rb->data_start++;
   } else {
      rb->count++;
   }
}

static int ringbuffer_pop(ring_buffer* rb) {
   int data = *rb->data_start;
   rb->data_start++;
   if (rb->data_start == rb->buffer_end)
      rb->data_start = rb->buffer;
   rb->count--;

   return data;
}

static int ringbuffer_full(ring_buffer* rb) {
   return rb->count == rb->size;
}

static int ringbuffer_empty(ring_buffer* rb) {
   return rb->count == 0;
}

/**********************************************************************/