afflib_pages.cpp

sleuthit-2.09 一个磁盘的工具集

}


static bool is_buffer_zero(unsigned char *buf,int buflen)
{
    for(int i=0;i<buflen;i++){
	if(buf[i]) return false;
    }
    return true;
}

/* Write a actual data segment to the disk */
int af_update_page(AFFILE *af,int64 pagenum,unsigned char *data,int datalen)
{
    char segname_buf[32];

#ifdef HAVE_OPENSSL_MD5_H
    /* Write out the hash of the page */
    unsigned char md5_buf[16];
    MD5(data,datalen,md5_buf);
    snprintf(segname_buf,sizeof(segname_buf),AF_PAGE_MD5,pagenum);
    af_update_seg(af,segname_buf,0,md5_buf,sizeof(md5_buf)); // ignore failure
#endif

    /* Check for bypass */
    if(af->v->write){
	int r = (*af->v->write)(af,data,af->image_pagesize * pagenum,datalen);
	if(r!=datalen) return -1;
	return 0;
    }


    struct affcallback_info acbi;
    int ret = 0;
    uint64 starting_pages_written = af->pages_written;

    /* Setup the callback structure */
    memset(&acbi,0,sizeof(acbi));
    acbi.info_version = 1;
    acbi.af = af->parent ? af->parent : af;
    acbi.pagenum = pagenum;
    acbi.bytes_to_write = datalen;

    /* Set up the segnment name */
    snprintf(segname_buf,sizeof(segname_buf),AF_PAGE,pagenum);
    size_t destLen = af->image_pagesize;	// it could be this big.

    /* Compress and write the data, if we are allowed to compress */
    if(af->compression_type != AF_COMPRESSION_ALG_NONE){
	unsigned char *cdata = (unsigned char *)malloc(destLen); // compressed data
	unsigned long *ldata = (unsigned long *)cdata; // allows me to reference as a buffer of unsigned longs
	if(cdata!=0){		// If data could be allocated

	    int cres = -1;		// compression results
	    unsigned int flag = 0;	// flag for data segment
	    int dont_compress = 0;

	    /* Try zero compression first; it's the best algorithm we have  */
	    if(is_buffer_zero(data,datalen)){
		acbi.compression_alg   = AF_PAGE_COMP_ALG_ZERO; 
		acbi.compression_level = AF_COMPRESSION_MAX;

		if(af->w_callback) { acbi.phase = 1; (*af->w_callback)(&acbi); }

		*ldata = htonl(datalen); // store the data length
		destLen = 4;		 // 4 bytes
		flag = AF_PAGE_COMPRESSED | AF_PAGE_COMP_ALG_ZERO | AF_PAGE_COMP_MAX; 
		cres = 0;

		acbi.compressed = 1;		// it was compressed
		if(af->w_callback) {acbi.phase = 2;(*af->w_callback)(&acbi);}
	    }

#ifdef USE_LZMA
	    if(cres!=0 && af->compression_type==AF_COMPRESSION_ALG_LZMA){ // try to compress with LZMA
		acbi.compression_alg   = AF_PAGE_COMP_ALG_LZMA;
		acbi.compression_level = 7; // right now, this is the level we use
		if(af->w_callback) { acbi.phase = 1; (*af->w_callback)(&acbi); }

		cres = lzma_compress(cdata,&destLen,data,datalen,9);
#if 0
		switch(cres){
		case 0:break;		// OKAY
		case 1: (*af->error_reporter)("LZMA: Unspecified Error\n");break;
		case 2: (*af->error_reporter)("LZMA: Memory Allocating Error\n");break;
		case 3: (*af->error_reporter)("LZMA: Output buffer OVERFLOW\n"); break;
		default: (*af->error_reporter)("LZMA: Unknown error %d\n",cres);break;
		}
#endif
		if(cres==0){
		    flag = AF_PAGE_COMPRESSED | AF_PAGE_COMP_ALG_LZMA;
		    acbi.compressed = 1;
		    if(af->w_callback) {acbi.phase = 2;(*af->w_callback)(&acbi);}
		}
		else {
		    /* Don't bother reporting LZMA errors; we just won't compress */
		    dont_compress = 1;
		    if(af->w_callback) {acbi.phase = 2;(*af->w_callback)(&acbi);}
		}
	    }
#endif

	    if(cres!=0
	       && af->compression_type==AF_COMPRESSION_ALG_ZLIB
	       && dont_compress==0){ // try to compress with zlib
		acbi.compression_alg   = AF_PAGE_COMP_ALG_ZLIB; // only one that we support
		acbi.compression_level = af->compression_level;
		if(af->w_callback) { acbi.phase = 1; (*af->w_callback)(&acbi); }

		cres = compress2((Bytef *)cdata, (uLongf *)&destLen,
				 (Bytef *)data,datalen, af->compression_level);

		if(cres==0){
		    flag = AF_PAGE_COMPRESSED | AF_PAGE_COMP_ALG_ZLIB;
		    if(af->compression_level == AF_COMPRESSION_MAX){
			flag |= AF_PAGE_COMP_MAX; // useful to know it can't be better
		    }
		}
		acbi.compressed = 1;	// it was compressed (or not compressed)
		if(af->w_callback) {acbi.phase = 2;(*af->w_callback)(&acbi);}
	    }

	    if(cres==0 && destLen < af->image_pagesize){
		/* Prepare to write out the compressed segment with compression */
		if(af->w_callback) {acbi.phase = 3;(*af->w_callback)(&acbi);}
		ret = af_update_seg(af,segname_buf,flag,cdata,destLen);
		acbi.bytes_written = destLen;
		if(af->w_callback) {acbi.phase = 4;(*af->w_callback)(&acbi);}
		if(ret==0){
		    af->pages_written++;
		    af->pages_compressed++;
		}
	    }
	    free(cdata);
	    cdata = 0;
	}
    }
    
    /* If a compressed segment was not written, write it uncompressed */
    if(af->pages_written == starting_pages_written){ 
	if(af->w_callback) {acbi.phase = 3;(*af->w_callback)(&acbi);}
	ret = af_update_seg(af,segname_buf,0,data,datalen);
	acbi.bytes_written = datalen;
	if(af->w_callback) {acbi.phase = 4;(*af->w_callback)(&acbi);} 
	if(ret==0){
	    acbi.bytes_written = datalen;	// because that is how much we wrote
	    af->pages_written++;
	}
    }
    return ret;
}

/****************************************************************
 *** Cache interface
 ****************************************************************/

/* The page cache is a read/write cache.
 * 
 * Pages that are read are cached after they are decompressed.
 * When new pages are fetched, we check the cache first to see if they are there; 
 * if so, they are satsfied by the cache.
 * 
 * Modifications are written to the cache, then dumped to the disk. 
 * 
 * The cache is managed by two functions:
 * af_cache_flush(af) - (prevously af_purge) 
 *      - Makes sure that all dirty buffers are written.
 *      - Sets af->pb=NULL (no current page)
 *      - (returns 0 if success, -1 if failure.)
 *
 * af_cache_writethrough(af,page,buf,buflen)
 *      - used for write bypass
 *
 * af_cache_load(af,page) - 
 *      - If page is already in the cache, return it.
 *      - If cache is filled, randomly discard a page
 *      - if page is on the disk, load the page.
 *      - Sets af->bp to be the page that was loaded.
 *
 */

static int cachetime = 0;


int af_cache_flush(AFFILE *af)
{
    if(af_trace) fprintf(af_trace,"af_cache_flush()\n");
    int ret = 0;
    for(int i=0;i<af->num_pbufs;i++){
	struct aff_pagebuf *p = &af->pbcache[i];
	if(p->pagebuf_valid && p->pagebuf_dirty){
	    if(af_update_page(af,p->pagenum,p->pagebuf,p->pagebuf_bytes)){
		ret = -1;		// got an error; keep going, though
	    }
	    p->pagebuf_dirty = 0;
	    if(af_trace) fprintf(af_trace,"af_cache_flush: slot %d page %qd flushed.\n",i,p->pagenum);
	}
    }
    return ret;				// now return the error that I might have gotten
}

/* If the page being written is in the cache, update it.
 * Question: would it make sense to copy the data anyway? I don't think so, because
 * the main use of writethrough is when imaging, and in that event you probably don't
 * want the extra memcpy.
 */
void af_cache_writethrough(AFFILE *af,int64 pagenum,const unsigned char *buf,int bufflen)
{
    for(int i=0;i<af->num_pbufs;i++){
	struct aff_pagebuf *p = &af->pbcache[i];
	if(p->pagenum_valid && p->pagenum == pagenum){
	    if(p->pagebuf_dirty){
		(*af->error_reporter)("af_cache_writethrough: overwriting page %"I64u".\n",pagenum);
		exit(-1);		// this shouldn't happen
	    }
	    memcpy(p->pagebuf,buf,bufflen);
	    memset(p->pagebuf+bufflen,0,af->image_pagesize-bufflen); // zero fill the rest
	    af->bytes_memcpy += bufflen;
	    p->pagebuf_valid = 1;	// we have a copy of it now.
	    p->pagebuf_dirty = 0;	// but it isn't dirty
	    p->last = cachetime++;
	}
    }
}
	
#ifdef HAVE_MALLOC_H
#include <malloc.h>
#endif

#ifndef HAVE_VALLOC
#define valloc malloc
#endif

struct aff_pagebuf *af_cache_alloc(AFFILE *af,int64 pagenum)
{
    if(af_trace) fprintf(af_trace,"af_cache_alloc(%p,%"I64d")\n",af,pagenum);
    af_cache_flush(af);			// make sure nothing in cache is dirty
    /* See if this page is already in the cache */
    for(int i=0;i<af->num_pbufs;i++){
	struct aff_pagebuf *p = &af->pbcache[i];
	if(p->pagenum_valid && p->pagenum==pagenum){
	    af->cache_hits++;
	    if(af_trace) fprintf(af_trace,"  page %"I64d" satisfied fromcache\n",pagenum);
	    p->last = cachetime++;
	    return p;
	}
    }

    af->cache_misses++;
    int slot = -1;
    /* See if there is an empty slot in the cache */
    for(int i=0;i<af->num_pbufs;i++){
	struct aff_pagebuf *p = &af->pbcache[i];
	if(p->pagenum_valid==0){
	    slot = i;
	    if(af_trace) fprintf(af_trace,"  slot %d given to page %"I64d"\n",slot,pagenum);
	    break;
	}
    }
    if(slot==-1){			
	/* Find the oldest cache entry */
	int oldest_i = 0;
	int oldest_t = af->pbcache[0].last;
	for(int i=1;i<af->num_pbufs;i++){
	    if(af->pbcache[i].last < oldest_t){
		oldest_t = af->pbcache[i].last;
		oldest_i = i;
	    }
	}
	slot = oldest_i;
	if(af_trace) fprintf(af_trace,"  slot %d assigned to page %"I64d"\n",slot,pagenum);
    }
    /* take over this slot */
    struct aff_pagebuf *p = &af->pbcache[slot];
    if(p->pagebuf==0){
	p->pagebuf = (unsigned char *)valloc(af->image_pagesize); // allocate to a page boundary
	if(p->pagebuf==0){
	    /* Malloc failed; See if we can just use the first slot */
	    slot = 0;
	    if(af->pbcache[0].pagebuf==0) return 0; // ugh. Cannot malloc?

	    /* First slot is available. Just use it. */
	    p = &af->pbcache[0];
	}
    }
    memset(p->pagebuf,0,af->image_pagesize); // clean object reuse
    p->pagenum = pagenum;
    p->pagenum_valid = 1;
    p->pagebuf_valid = 0;
    p->pagebuf_dirty = 0;
    p->last = cachetime++;
    if(af_trace){
	fprintf(af_trace,"   current pages in cache: ");
	for(int i=0;i<af->num_pbufs;i++){
	    fprintf(af_trace," %"I64d,af->pbcache[i].pagenum);
	}
	fprintf(af_trace,"\n");
    }
    return p;
}